feat(stream): allow Stream/Flow's stage() to add multiple stages

zhutmost · Jul 9, 2024 · 84f8ee1 · 84f8ee1
1 parent ffe757d
commit 84f8ee1
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 2 deletions.
diff --git a/chipmunk/src/stream/Flow.scala b/chipmunk/src/stream/Flow.scala
@@ -124,7 +124,14 @@ class FlowIO[T <: Data](gen: T) extends Valid[T](gen) with IsMasterSlave {
   }
 
   /** Alias of [[pipeForward]]. */
-  def stage(): FlowIO[T] = pipeForward()
+  def stage(n: Int = 1): FlowIO[T] = {
+    require(n >= 0, "The stage number must be greater than or equal to 0.")
+    var ret = this
+    for (_ <- 0 until n) {
+      ret = ret.pipeForward()
+    }
+    ret
+  }
 
   /** Throw transactions when `cond` is False. */
   def takeWhen(cond: Bool): FlowIO[T] = {

diff --git a/chipmunk/src/stream/Stream.scala b/chipmunk/src/stream/Stream.scala
@@ -257,7 +257,14 @@ class StreamIO[T <: Data](gen: T) extends DecoupledIO[T](gen) with IsMasterSlave
   }
 
   /** Alias of [[pipeForward]]. */
-  def stage(): StreamIO[T] = pipeForward()
+  def stage(n: Int = 1): StreamIO[T] = {
+    require(n >= 0, "The stage number must be greater than or equal to 0.")
+    var ret = this
+    for (_ <- 0 until n) {
+      ret = ret.pipeForward()
+    }
+    ret
+  }
 
   /** Block this stream when `cond` is False. */
   def continueWhen(cond: => Bool): StreamIO[T] = {