From 7bdd9da49d5de9b19a7cf112fff902cdb9bf4962 Mon Sep 17 00:00:00 2001
From: Mike Urbach <mikeurbach@gmail.com>
Date: Fri, 24 Dec 2021 10:42:40 -0700
Subject: [PATCH] [StaticLogic] Add a helper to get a pipeline stage number.

The stage number is computed as its index in the stages block. This is
useful to passes that may want to give a stage a meaningful name.
---
 include/circt/Dialect/StaticLogic/StaticLogic.td |  1 +
 lib/Dialect/StaticLogic/StaticLogicOps.cpp       | 12 ++++++++++++
 2 files changed, 13 insertions(+)

diff --git a/include/circt/Dialect/StaticLogic/StaticLogic.td b/include/circt/Dialect/StaticLogic/StaticLogic.td
index 2b749f0802fe..1e89510ed455 100644
--- a/include/circt/Dialect/StaticLogic/StaticLogic.td
+++ b/include/circt/Dialect/StaticLogic/StaticLogic.td
@@ -199,6 +199,7 @@ def PipelineStageOp : Op<StaticLogic_Dialect, "pipeline.stage",
 
   let extraClassDeclaration = [{
     Block &getBodyBlock() { return body().front(); }
+    unsigned getStageNumber();
   }];
 }
 
diff --git a/lib/Dialect/StaticLogic/StaticLogicOps.cpp b/lib/Dialect/StaticLogic/StaticLogicOps.cpp
index 73ee728d3f20..6225132e22da 100644
--- a/lib/Dialect/StaticLogic/StaticLogicOps.cpp
+++ b/lib/Dialect/StaticLogic/StaticLogicOps.cpp
@@ -208,6 +208,18 @@ static LogicalResult verifyPipelineRegisterOp(PipelineRegisterOp op) {
   return success();
 }
 
+unsigned PipelineStageOp::getStageNumber() {
+  unsigned number = 0;
+  auto *op = getOperation();
+  auto parent = op->getParentOfType<PipelineWhileOp>();
+  Operation *stage = &parent.getStagesBlock().front();
+  while (stage != op && stage->getNextNode()) {
+    ++number;
+    stage = stage->getNextNode();
+  }
+  return number;
+}
+
 //===----------------------------------------------------------------------===//
 // PipelineTerminatorOp
 //===----------------------------------------------------------------------===//