Implement DecoupledIO.map utility (#2646) (#2649)

(cherry picked from commit b20df1d6cda03f6eef28ee480e0aade914c5face)

Co-authored-by: Jared Barocsi <82000041+jared-barocsi@users.noreply.github.com>

2 files changed, 107 insertions, 1 deletions

diff --git a/src/main/scala/chisel3/util/Decoupled.scala b/src/main/scala/chisel3/util/Decoupled.scala
index f02a4116..b21bd04f 100644
--- a/src/main/scala/chisel3/util/Decoupled.scala
+++ b/src/main/scala/chisel3/util/Decoupled.scala
@@ -99,7 +99,22 @@ object ReadyValidIO {
   * of ready or valid.
   * @param gen the type of data to be wrapped in DecoupledIO
   */
-class DecoupledIO[+T <: Data](gen: T) extends ReadyValidIO[T](gen)
+class DecoupledIO[+T <: Data](gen: T) extends ReadyValidIO[T](gen) {
+
+  /** Applies the supplied functor to the bits of this interface, returning a new
+    * typed DecoupledIO interface.
+    * @param f The function to apply to this DecoupledIO's 'bits' with return type B
+    * @return a new DecoupledIO of type B
+    */
+  def map[B <: Data](f: T => B): DecoupledIO[B] = {
+    val _map_bits = f(bits)
+    val _map = Wire(Decoupled(chiselTypeOf(_map_bits)))
+    _map.bits := _map_bits
+    _map.valid := valid
+    ready := _map.ready
+    _map
+  }
+}
 
 /** This factory adds a decoupled handshaking protocol to a data bundle. */
 object Decoupled {
diff --git a/src/test/scala/chiselTests/DecoupledSpec.scala b/src/test/scala/chiselTests/DecoupledSpec.scala
index 2d305f4a..69d74aab 100644
--- a/src/test/scala/chiselTests/DecoupledSpec.scala
+++ b/src/test/scala/chiselTests/DecoupledSpec.scala
@@ -17,4 +17,95 @@ class DecoupledSpec extends ChiselFlatSpec {
       assert(io.asUInt.widthOption.get === 4)
     })
   }
+
+  "Decoupled.map" should "apply a function to a wrapped Data" in {
+    val chirrtl = ChiselStage
+      .emitChirrtl(new Module {
+        val enq = IO(Flipped(Decoupled(UInt(8.W))))
+        val deq = IO(Decoupled(UInt(8.W)))
+        deq <> enq.map(_ + 1.U)
+      })
+
+    // Check for data assignment
+    chirrtl should include("""node _deq_map_bits_T = add(enq.bits, UInt<1>("h1")""")
+    chirrtl should include("""node _deq_map_bits = tail(_deq_map_bits_T, 1)""")
+    chirrtl should include("""_deq_map.bits <= _deq_map_bits""")
+    chirrtl should include("""deq <= _deq_map""")
+
+    // Check for back-pressure (ready signal is driven in the opposite direction of bits + valid)
+    chirrtl should include("""enq.ready <= _deq_map.ready""")
+  }
+
+  "Decoupled.map" should "apply a function to a wrapped Bundle" in {
+    class TestBundle extends Bundle {
+      val foo = UInt(8.W)
+      val bar = UInt(8.W)
+      val fizz = Bool()
+      val buzz = Bool()
+    }
+
+    // Add one to foo, subtract one from bar, set fizz to false and buzz to true
+    def func(t: TestBundle): TestBundle = {
+      val res = Wire(new TestBundle)
+
+      res.foo := t.foo + 1.U
+      res.bar := t.bar - 1.U
+      res.fizz := false.B
+      res.buzz := true.B
+
+      res
+    }
+
+    val chirrtl = ChiselStage
+      .emitChirrtl(new Module {
+        val enq = IO(Flipped(Decoupled(new TestBundle)))
+        val deq = IO(Decoupled(new TestBundle))
+        deq <> enq.map(func)
+      })
+
+    // Check for data assignment
+    chirrtl should include("""wire _deq_map_bits : { foo : UInt<8>, bar : UInt<8>, fizz : UInt<1>, buzz : UInt<1>}""")
+
+    chirrtl should include("""node _deq_map_bits_res_foo_T = add(enq.bits.foo, UInt<1>("h1")""")
+    chirrtl should include("""node _deq_map_bits_res_foo_T_1 = tail(_deq_map_bits_res_foo_T, 1)""")
+    chirrtl should include("""_deq_map_bits.foo <= _deq_map_bits_res_foo_T_1""")
+
+    chirrtl should include("""node _deq_map_bits_res_bar_T = sub(enq.bits.bar, UInt<1>("h1")""")
+    chirrtl should include("""node _deq_map_bits_res_bar_T_1 = tail(_deq_map_bits_res_bar_T, 1)""")
+    chirrtl should include("""_deq_map_bits.bar <= _deq_map_bits_res_bar_T_1""")
+
+    chirrtl should include("""_deq_map_bits.fizz <= UInt<1>("h0")""")
+    chirrtl should include("""_deq_map_bits.buzz <= UInt<1>("h1")""")
+
+    chirrtl should include("""_deq_map.bits <= _deq_map_bits""")
+    chirrtl should include("""deq <= _deq_map""")
+
+    // Check for back-pressure (ready signal is driven in the opposite direction of bits + valid)
+    chirrtl should include("""enq.ready <= _deq_map.ready""")
+  }
+
+  "Decoupled.map" should "apply a function to a wrapped Bundle and return a different typed DecoupledIO" in {
+    class TestBundle extends Bundle {
+      val foo = UInt(8.W)
+      val bar = UInt(8.W)
+    }
+
+    val chirrtl = ChiselStage
+      .emitChirrtl(new Module {
+        val enq = IO(Flipped(Decoupled(new TestBundle)))
+        val deq = IO(Decoupled(UInt(8.W)))
+        deq <> enq.map(bundle => bundle.foo & bundle.bar)
+      })
+
+    // Check that the _map wire wraps a UInt and not a TestBundle
+    chirrtl should include("""wire _deq_map : { flip ready : UInt<1>, valid : UInt<1>, bits : UInt<8>}""")
+
+    // Check for data assignment
+    chirrtl should include("""node _deq_map_bits = and(enq.bits.foo, enq.bits.bar)""")
+    chirrtl should include("""_deq_map.bits <= _deq_map_bits""")
+    chirrtl should include("""deq <= _deq_map""")
+
+    // Check for back-pressure (ready signal is driven in the opposite direction of bits + valid)
+    chirrtl should include("""enq.ready <= _deq_map.ready""")
+  }
 }