reduceTree() now operates on Seq (#2292)

preserves input/output information of the type being reduced.

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

2 files changed, 37 insertions, 3 deletions

diff --git a/core/src/main/scala/chisel3/Aggregate.scala b/core/src/main/scala/chisel3/Aggregate.scala
index dbde7068..600e2d11 100644
--- a/core/src/main/scala/chisel3/Aggregate.scala
+++ b/core/src/main/scala/chisel3/Aggregate.scala
@@ -338,11 +338,11 @@ sealed class Vec[T <: Data] private[chisel3] (gen: => T, val length: Int)
   def do_reduceTree(redOp: (T, T) => T, layerOp: (T) => T = (x: T) => x)
                    (implicit sourceInfo: SourceInfo, compileOptions: CompileOptions) : T = {
     require(!isEmpty, "Cannot apply reduction on a vec of size 0")
-    var curLayer = this
+    var curLayer : Seq[T] = this
     while (curLayer.length > 1) {
-      curLayer = VecInit(curLayer.grouped(2).map( x =>
+      curLayer = curLayer.grouped(2).map( x =>
         if (x.length == 1) layerOp(x(0)) else redOp(x(0), x(1))
-      ).toSeq)
+      ).toSeq
     }
     curLayer(0)
   }
diff --git a/src/test/scala/chiselTests/Vec.scala b/src/test/scala/chiselTests/Vec.scala
index 97aea909..24ba0bf8 100644
--- a/src/test/scala/chiselTests/Vec.scala
+++ b/src/test/scala/chiselTests/Vec.scala
@@ -313,6 +313,36 @@ class ModuleIODynamicIndexTester(n: Int) extends BasicTester {
   when (done) { stop() }
 }
 
+class ReduceTreeTester() extends BasicTester {
+  class FooIO[T <: Data](n: Int, private val gen: T) extends Bundle {
+    val in = Flipped(Vec(n, new DecoupledIO(gen)))
+    val out = new DecoupledIO(gen)
+  }
+
+  class Foo[T <: Data](n: Int, private val gen: T) extends Module {
+    val io = IO(new FooIO(n, gen))
+
+    def foo(a: DecoupledIO[T], b: DecoupledIO[T]) = {
+      a.ready := true.B
+      b.ready := true.B
+      val out = Wire(new DecoupledIO(gen))
+
+      out.valid := true.B
+
+      val regSel = RegInit(false.B)
+      out.bits := Mux(regSel, a.bits, b.bits)
+      out.ready := a.ready
+      out
+    }
+
+    io.out <> io.in.reduceTree(foo)
+  }
+
+  val dut = Module(new Foo(5, UInt(5.W)))
+  dut.io := DontCare
+  stop()
+}
+
 class VecSpec extends ChiselPropSpec with Utils {
   // Disable shrinking on error.
   implicit val noShrinkListVal = Shrink[List[Int]](_ => Stream.empty)
@@ -456,4 +486,8 @@ class VecSpec extends ChiselPropSpec with Utils {
         }}
     }
   }
+
+  property("reduceTree should preserve input/output type") {
+      assertTesterPasses { new ReduceTreeTester() }
+  }
 }