From 186fc15f8fde0a7921c04d229c01c67335e22d31 Mon Sep 17 00:00:00 2001
From: Stan Lee <slee93@stanford.edu>
Date: Mon, 19 Jan 2026 12:10:48 -0800
Subject: [PATCH] passes simple test

---
 passes/silimate/reg_rename.cc | 79 +++++++++++++++++++++++++++++++----
 1 file changed, 70 insertions(+), 9 deletions(-)

diff --git a/passes/silimate/reg_rename.cc b/passes/silimate/reg_rename.cc
index a22aa06cd..9baf36321 100644
--- a/passes/silimate/reg_rename.cc
+++ b/passes/silimate/reg_rename.cc
@@ -24,6 +24,17 @@
 USING_YOSYS_NAMESPACE
 PRIVATE_NAMESPACE_BEGIN
 
+struct CellTracker {
+	Cell *cell;
+	int index;
+};
+
+struct RegTracker {
+	std::map<std::string, CellTracker> renamedRegs;
+	std::string origRegWidth;
+	Module *module;
+};
+
 struct RegRenamePass : public Pass {
 	RegRenamePass() : Pass("reg_rename", "renames register output wires to the correct register name") { }
 	void help() override
@@ -49,7 +60,7 @@ struct RegRenamePass : public Pass {
 
 		// Data structure used to keep track of multi-bit registers.
 		// Relevant for correct register annotation.
-
+		std::map<std::string, RegTracker> regTrackers;
 
 		// Regex to match register output wires
 		// .*_reg[NUMBER] or .*_reg, can match NUMBER and part before _reg
@@ -62,23 +73,23 @@ struct RegRenamePass : public Pass {
 				std::string name = cell->name.c_str();
 				if (std::regex_match(name, match, reg_regex)) {
 
-
-					std::string origRegWidth = cell->get_string_attribute("$ORIG_REG_WIDTH");
-					log("Original register width for cell %s: %s\n", cell->name.c_str(), origRegWidth.c_str());
-
 					// baseName is the part before _reg
 					std::string baseName = match[1].str();
 					std::string registerName = baseName;
-					if (match.size() > 2 && match[2].matched) {
+
+					// Check if the register is a multi-bit register
+					bool isMultiBit = match.size() > 2 && match[2].matched;
+					std::string indexStr;
+					if (isMultiBit) {
 						// indexStr is the NUMBER in .*_reg[NUMBER]
-						std::string indexStr = match[2].str();
-						registerName += indexStr;
+						indexStr = match[2].str();
+						registerName += "_" + indexStr;
 					}
 
 					for (auto conn : cell->connections()) {
 						if (conn.first == ID::Q && conn.second.is_wire()) {
 							Wire *wire = conn.second.as_wire();
-							
+
 							// Skip if this wire is a module port (input/output)
 							if (wire->port_input || wire->port_output) {
 								log("Skipping port wire %s in register renaming for cell %s in module %s\n", 
@@ -94,6 +105,15 @@ struct RegRenamePass : public Pass {
 							// Rename register
 							log("Renaming register wire %s to %s for cell %s in module %s\n", 
 								wire->name.c_str(), registerName, log_id(cell), log_id(module));
+
+							// Log relevant information for multi-bit registers for wire reconstruction
+							if (isMultiBit) {
+								std::string origRegWidth = cell->get_string_attribute("$ORIG_REG_WIDTH");
+								regTrackers[baseName].origRegWidth = origRegWidth;
+								regTrackers[baseName].renamedRegs[registerName] = CellTracker{cell, std::stoi(indexStr)};
+								regTrackers[baseName].module = module;
+							}
+								
 							module->rename(wire, registerName);
 							count++;
 						}
@@ -101,6 +121,47 @@ struct RegRenamePass : public Pass {
 				}
 			}
 		}
+
+	for (const auto &[baseName, regTracker] : regTrackers) {
+
+		// Get the module for this register
+		Module *mod = regTracker.module;
+
+		// Create a new wire for the multi-bit register
+		int width = std::stoi(regTracker.origRegWidth);
+		log("Creating new wire %s for register %s with width %d\n", baseName.c_str(), baseName.c_str(), width);
+		Wire *newWire = mod->addWire(baseName, width);
+
+		pool<Wire *> oldWires;
+
+		// Connect the renamed registers to the corresponding index of the new wire
+		for (const auto &[renamedRegName, cellTracker] : regTracker.renamedRegs) {
+
+			// Get the old wire (the Q output that was renamed)
+			Wire *oldWire = cellTracker.cell->getPort(ID::Q).as_wire();
+
+			// Get the index of the renamed register
+			int index = cellTracker.index;
+			log("Connecting renamed register %s to index %d\n", renamedRegName.c_str(), index);
+
+			// Connect the renamed register to the corresponding index of the new wire
+			mod->connect(SigSpec(newWire, index, 1), cellTracker.cell->getPort(ID::Q));
+
+			// Replace all uses of oldWire with newWire[index] throughout the module
+			auto rewriter = [&](SigSpec &sig) {
+				sig.replace(SigBit(oldWire), SigSpec(newWire, index, 1));
+			};
+			mod->rewrite_sigspecs(rewriter);
+
+			// Add the old wire to the list of old wires to delete
+			oldWires.insert(oldWire);
+		}
+
+		// Delete the old wires
+		mod->remove(oldWires);
+	}
+
+
 		log("Renamed %d registers in %d modules\n", count, moduleCount);
 		log_flush();
 	}