From 454cd8a9385629f20b3125120debd1acab1fc60b Mon Sep 17 00:00:00 2001
From: ducky <richard.lin@berkeley.edu>
Date: Thu, 14 Nov 2019 13:57:43 -0800
Subject: [PATCH] Support Bundle literal peek

---
 src/main/scala/chisel3/tester/package.scala   |  7 +++++++
 .../chisel3/tests/BundleLiteralsSpec.scala    | 19 +++++++++++--------
 2 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/src/main/scala/chisel3/tester/package.scala b/src/main/scala/chisel3/tester/package.scala
index 59346cd96..5c500bda5 100644
--- a/src/main/scala/chisel3/tester/package.scala
+++ b/src/main/scala/chisel3/tester/package.scala
@@ -5,6 +5,7 @@ package chisel3
 import scala.language.implicitConversions
 import chisel3.tester.internal._
 import chisel3.experimental.{DataMirror, Direction, FixedPoint}
+import chisel3.experimental.BundleLiterals._
 import chisel3.util._
 
 class NotLiteralException(message: String) extends Exception(message)
@@ -61,6 +62,12 @@ package object tester {
         val multiplier = math.pow(2, x.binaryPoint.get)
         (Context().backend.peekBits(x, stale).toDouble / multiplier).F(x.binaryPoint).asInstanceOf[T]
       }
+      case (x: Bundle) => {
+        val elementValueFns = x.elements.map { case (name: String, elt: Data) =>
+          (y: Bundle) => (y.elements(name), elt.peekWithStale(stale))
+        }.toSeq
+        chiselTypeOf(x).Lit(elementValueFns: _*).asInstanceOf[T]
+      }
       case x => throw new LiteralTypeException(s"don't know how to peek $x")
     }
 
diff --git a/src/test/scala/chisel3/tests/BundleLiteralsSpec.scala b/src/test/scala/chisel3/tests/BundleLiteralsSpec.scala
index 9acc8ec97..c46df2af9 100644
--- a/src/test/scala/chisel3/tests/BundleLiteralsSpec.scala
+++ b/src/test/scala/chisel3/tests/BundleLiteralsSpec.scala
@@ -53,17 +53,20 @@ class BundleLiteralsSpec extends FlatSpec with ChiselScalatestTester with Matche
     }
   }
 
-  // peek on bundle not supported yet, this test will fail and should
-  // be altered when BundleLiteral peeking works
-  // it is not altogether what the use case is for peeking a bundle
-  // possibly to poke that value somewhere else
-  // this should be considered when peeking Bundles is supported
-  ignore should "return a BundleLiteral when peeking" in {
+  it should "return a Bundle literal when peeking" in {
     test(new PassthroughModule(new DoubleElements)) { c =>
       c.in.poke(chiselTypeOf(c.in).Lit(_.a -> 0.U, _.b -> 1.U))
       val output = c.out.peek()
-      output.a === 0.U should be(true.B)
-      output.a === 1.U should be(true.B)
+      output.a.litValue should be (0)
+      output.b.litValue should be (1)
+    }
+  }
+
+  it should "roundtrip Bundle literals" in {
+    test(new PassthroughModule(new DoubleElements)) { c =>
+      c.in.poke(chiselTypeOf(c.in).Lit(_.a -> 0.U, _.b -> 1.U))
+      c.in.poke(c.out.peek())
+      c.out.expect(chiselTypeOf(c.in).Lit(_.a -> 0.U, _.b -> 1.U))
     }
   }
 }