Added VecInit factory methods (fill,iterate) (#2059)

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

3 files changed, 58 insertions, 3 deletions

diff --git a/core/src/main/scala/chisel3/Aggregate.scala b/core/src/main/scala/chisel3/Aggregate.scala
index 45ecec36..3a766628 100644
--- a/core/src/main/scala/chisel3/Aggregate.scala
+++ b/core/src/main/scala/chisel3/Aggregate.scala
@@ -590,6 +590,33 @@ object VecInit extends SourceInfoDoc {
   /** @group SourceInfoTransformMacro */
   def do_tabulate[T <: Data](n: Int)(gen: (Int) => T)(implicit sourceInfo: SourceInfo, compileOptions: CompileOptions): Vec[T] =
     apply((0 until n).map(i => gen(i)))
+
+  /** Creates a new [[Vec]] of length `n` composed of the result of the given
+   * function applied to an element of data type T.
+   * 
+   * @param n number of elements in the vector
+   * @param gen function that takes in an element T and returns an output 
+   * element of the same type
+   */
+  def fill[T <: Data](n: Int)(gen: => T): Vec[T] = macro VecTransform.fill
+
+  /** @group SourceInfoTransformMacro */
+  def do_fill[T <: Data](n: Int)(gen: => T)(implicit sourceInfo: SourceInfo, compileOptions: CompileOptions): Vec[T] =
+    apply(Seq.fill(n)(gen))
+
+  /** Creates a new [[Vec]] of length `n` composed of the result of the given
+   * function applied to an element of data type T.
+   * 
+   * @param start First element in the Vec
+   * @param len Lenth of elements in the Vec
+   * @param f Function that applies the element T from previous index and returns the output 
+   * element to the next index
+   */
+  def iterate[T <: Data](start: T, len: Int)(f: (T) => T): Vec[T] = macro VecTransform.iterate
+  
+  /** @group SourceInfoTransformMacro */
+  def do_iterate[T <: Data](start: T, len: Int)(f: (T) => T)(implicit sourceInfo: SourceInfo, compileOptions: CompileOptions): Vec[T] = 
+    apply(Seq.iterate(start, len)(f)) 
 }
 
 /** A trait for [[Vec]]s containing common hardware generators for collection
@@ -964,7 +991,7 @@ abstract class Bundle(implicit compileOptions: CompileOptions) extends Record {
       getBundleField(m) match {
         case Some(d: Data) =>
           requireIsChiselType(d)
-
+          
           if (nameMap contains m.getName) {
             require(nameMap(m.getName) eq d)
           } else {
@@ -1268,3 +1295,4 @@ abstract class Bundle(implicit compileOptions: CompileOptions) extends Record {
     */
   override def toPrintable: Printable = toPrintableHelper(elements.toList.reverse)
 }
+
diff --git a/macros/src/main/scala/chisel3/internal/sourceinfo/SourceInfoTransform.scala b/macros/src/main/scala/chisel3/internal/sourceinfo/SourceInfoTransform.scala
index 4533aa39..d7c301e9 100644
--- a/macros/src/main/scala/chisel3/internal/sourceinfo/SourceInfoTransform.scala
+++ b/macros/src/main/scala/chisel3/internal/sourceinfo/SourceInfoTransform.scala
@@ -84,6 +84,9 @@ class VecTransform(val c: Context) extends SourceInfoTransformMacro {
   def fill(n: c.Tree)(gen: c.Tree): c.Tree = {
     q"$thisObj.do_fill($n)($gen)($implicitSourceInfo, $implicitCompileOptions)"
   }
+  def iterate(start: c.Tree, len: c.Tree)(f: c.Tree): c.Tree = {
+    q"$thisObj.do_iterate($start,$len)($f)($implicitSourceInfo, $implicitCompileOptions)"
+  }
   def contains(x: c.Tree)(ev: c.Tree): c.Tree = {
     q"$thisObj.do_contains($x)($implicitSourceInfo, $ev, $implicitCompileOptions)"
   }
diff --git a/src/test/scala/chiselTests/Vec.scala b/src/test/scala/chiselTests/Vec.scala
index 1327913d..f3160c2e 100644
--- a/src/test/scala/chiselTests/Vec.scala
+++ b/src/test/scala/chiselTests/Vec.scala
@@ -2,6 +2,8 @@
 
 package chiselTests
 
+import org.scalacheck._
+
 import chisel3._
 import chisel3.stage.ChiselStage
 import chisel3.testers.BasicTester
@@ -104,6 +106,21 @@ class TabulateTester(n: Int) extends BasicTester {
   stop()
 }
 
+class FillTester(n: Int, value: Int) extends BasicTester {
+  val x = VecInit(Array.fill(n)(value.U))
+  val u = VecInit.fill(n)(value.U)
+
+  assert(x.asUInt() === u.asUInt(), s"Expected Vec to be filled like $x, instead VecInit.fill created $u")
+  stop()
+}
+
+class IterateTester(start: Int, len: Int)(f: UInt => UInt) extends BasicTester {
+  val controlVec = VecInit(Seq.iterate(start.U, len)(f))
+  val testVec = VecInit.iterate(start.U, len)(f)
+  assert(controlVec.asUInt() === testVec.asUInt(), s"Expected Vec to be filled like $controlVec, instead creaeted $testVec\n")
+  stop()
+}
+
 class ShiftRegisterTester(n: Int) extends BasicTester {
   val (cnt, wrap) = Counter(true.B, n*2)
   val shifter = Reg(Vec(n, UInt((log2Ceil(n) max 1).W)))
@@ -160,7 +177,6 @@ class PassthroughModuleTester extends Module {
   assert(io.out === 123.U)
 }
 
-
 class ModuleIODynamicIndexTester(n: Int) extends BasicTester {
   val duts = VecInit(Seq.fill(n)(Module(new PassthroughModule).io))
   val tester = Module(new PassthroughModuleTester)
@@ -219,10 +235,18 @@ class VecSpec extends ChiselPropSpec with Utils {
     }
   }
 
-  property("Vecs should tabulate correctly") {
+  property("VecInit should tabulate correctly") {
     forAll(smallPosInts) { (n: Int) => assertTesterPasses{ new TabulateTester(n) } }
   }
 
+  property("VecInit should fill correctly") {
+    forAll(smallPosInts, Gen.choose(0, 50)) { (n: Int, value: Int) => assertTesterPasses{ new FillTester(n, value) } }
+  }
+
+  property("VecInit should iterate correctly") {
+    forAll(Gen.choose(1, 10), smallPosInts) { (start: Int, len: Int) => assertTesterPasses{ new IterateTester(start, len)(x => x + 50.U)}}
+  }
+
   property("Regs of vecs should be usable as shift registers") {
     forAll(smallPosInts) { (n: Int) => assertTesterPasses{ new ShiftRegisterTester(n) } }
   }