From 84f8ee18aacd0b412fe6b8c7b0f9b832caf91a75 Mon Sep 17 00:00:00 2001
From: zhutmost <zhutmost@outlook.com>
Date: Tue, 9 Jul 2024 13:20:51 +0800
Subject: [PATCH] feat(stream): allow Stream/Flow's stage() to add multiple
 stages

---
 chipmunk/src/stream/Flow.scala   | 9 ++++++++-
 chipmunk/src/stream/Stream.scala | 9 ++++++++-
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/chipmunk/src/stream/Flow.scala b/chipmunk/src/stream/Flow.scala
index 0cdb71d..6da4db6 100644
--- 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
index 07f1cac..3fda566 100644
--- 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] = {