complexassign test

2 files changed, 50 insertions, 0 deletions

diff --git a/src/test/scala/chiselTests/ChiselSpec.scala b/src/test/scala/chiselTests/ChiselSpec.scala
index 002ace8c..f3dad3dc 100644
--- a/src/test/scala/chiselTests/ChiselSpec.scala
+++ b/src/test/scala/chiselTests/ChiselSpec.scala
@@ -41,5 +41,6 @@ class ChiselPropSpec extends PropSpec with PropertyChecks {
   val safeUIntWidth = Gen.choose(1, 31) 
   val safeUInts = Gen.choose(0, (1 << 30))
   val vecSizes = Gen.choose(0, 4)
+  def enSequence(n: Int) = Gen.containerOfN[List,Boolean](n,Gen.oneOf(true,false))
 }
 
diff --git a/src/test/scala/chiselTests/ComplexAssign.scala b/src/test/scala/chiselTests/ComplexAssign.scala
new file mode 100644
index 00000000..ab940440
--- /dev/null
+++ b/src/test/scala/chiselTests/ComplexAssign.scala
@@ -0,0 +1,49 @@
+package chiselTests
+import Chisel._
+import org.scalatest._
+import org.scalatest.prop._
+import Chisel.testers.BasicTester
+
+class Complex[T <: Data](val re: T, val im: T) extends Bundle {
+  override def cloneType: this.type =
+    new Complex(re.cloneType, im.cloneType).asInstanceOf[this.type]
+}
+
+class ComplexAssign(w: Int) extends Module {
+  val io = new Bundle {
+    val e   = new Bool(INPUT)
+    val in  = new Complex(UInt(width = w), UInt(width = w)).asInput
+    val out = new Complex(UInt(width = w), UInt(width = w)).asOutput
+  }
+  when (io.e) {
+    val tmp = Wire(new Complex(UInt(width = w), UInt(width = w)))
+    tmp := io.in
+    io.out.re := tmp.re
+    io.out.im := tmp.im
+  } .otherwise {
+    io.out.re := UInt(0)
+    io.out.im := UInt(0)
+  }
+}
+
+class ComplexAssignSpec extends ChiselPropSpec {
+
+  class ComplexAssignTester(enList: List[Boolean], re: Int, im: Int) extends BasicTester {
+    val (cnt, wrap) = Counter(Bool(true), enList.size)
+    val dut = Module(new ComplexAssign(32))
+    dut.io.in.re := UInt(re)
+    dut.io.in.im := UInt(im)
+    dut.io.e := Vec(enList.map(Bool(_)))(cnt)
+    val re_correct = dut.io.out.re === Mux(dut.io.e, dut.io.in.re, UInt(0))
+    val im_correct = dut.io.out.im === Mux(dut.io.e, dut.io.in.im, UInt(0))
+    when(!re_correct || !im_correct) {
+      io.done := Bool(true); io.error := cnt 
+    } .elsewhen(wrap) { io.done := Bool(true) }
+  }
+     
+  property("All complex assignments should return the correct result") {
+    forAll(enSequence(16), safeUInts, safeUInts) { (en: List[Boolean], re: Int, im: Int) =>
+      assert(execute{ new ComplexAssignTester(en, re, im) }) 
+    }
+  }
+}