From 49b5331d78d4ff67350a6dde05816a2f3c10eb04 Mon Sep 17 00:00:00 2001
From: Schuyler Eldridge <schuyler.eldridge@sifive.com>
Date: Sat, 2 Mar 2024 19:50:44 -0500
Subject: [PATCH] [LowerToHW][nfc] Use function to record mapping

Add a member function for recording a mapping of old FIRRTL module to new
HW module.  Use this to avoid some tedious repeated code that manipulates
two mappings (old to new and new to old) directly.

Signed-off-by: Schuyler Eldridge <schuyler.eldridge@sifive.com>
---
 lib/Conversion/FIRRTLToHW/LowerToHW.cpp | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/lib/Conversion/FIRRTLToHW/LowerToHW.cpp b/lib/Conversion/FIRRTLToHW/LowerToHW.cpp
index 3e56fa7fde9c..799603b22b1b 100644
--- a/lib/Conversion/FIRRTLToHW/LowerToHW.cpp
+++ b/lib/Conversion/FIRRTLToHW/LowerToHW.cpp
@@ -273,6 +273,11 @@ struct CircuitLoweringState {
     return it != newToOldModuleMap.end() ? it->second : nullptr;
   }
 
+  void recordModuleMapping(Operation *oldFMod, Operation *newHWMod) {
+    oldToNewModuleMap[oldFMod] = newHWMod;
+    newToOldModuleMap[newHWMod] = oldFMod;
+  }
+
   // Process remaining annotations and emit warnings on unprocessed annotations
   // still remaining in the annoSet.
   void processRemainingAnnotations(Operation *op, const AnnotationSet &annoSet);
@@ -610,8 +615,7 @@ void FIRRTLModuleLowering::runOnOperation() {
               if (!loweredMod)
                 return failure();
 
-              state.oldToNewModuleMap[&op] = loweredMod;
-              state.newToOldModuleMap[loweredMod] = &op;
+              state.recordModuleMapping(&op, loweredMod);
               modulesToProcess.push_back(loweredMod);
               // Lower all the alias types.
               module.walk([&](Operation *op) {
@@ -628,8 +632,7 @@ void FIRRTLModuleLowering::runOnOperation() {
                   lowerExtModule(extModule, topLevelModule, state);
               if (!loweredMod)
                 return failure();
-              state.oldToNewModuleMap[&op] = loweredMod;
-              state.newToOldModuleMap[loweredMod] = &op;
+              state.recordModuleMapping(&op, loweredMod);
               return success();
             })
             .Case<FMemModuleOp>([&](auto memModule) {
@@ -637,8 +640,7 @@ void FIRRTLModuleLowering::runOnOperation() {
                   lowerMemModule(memModule, topLevelModule, state);
               if (!loweredMod)
                 return failure();
-              state.oldToNewModuleMap[&op] = loweredMod;
-              state.newToOldModuleMap[loweredMod] = &op;
+              state.recordModuleMapping(&op, loweredMod);
               return success();
             })
             .Default([&](Operation *op) {