From ce5cb8c06c0839cc16ed50db7e64d1ec4c05490f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vladim=C3=ADr=20=C5=A0till?= <git@vstill.eu>
Date: Mon, 27 May 2024 08:47:12 +0200
Subject: [PATCH] Generalization & minor refactoring in RenameMap (#4677)

---
 frontends/p4/uniqueNames.cpp | 12 ++++++------
 frontends/p4/uniqueNames.h   | 26 +++++++++++++++++++++-----
 2 files changed, 27 insertions(+), 11 deletions(-)

diff --git a/frontends/p4/uniqueNames.cpp b/frontends/p4/uniqueNames.cpp
index 8090fda4b72..f217cb744a7 100644
--- a/frontends/p4/uniqueNames.cpp
+++ b/frontends/p4/uniqueNames.cpp
@@ -21,12 +21,12 @@ limitations under the License.
 
 namespace P4 {
 
-void RenameMap::setNewName(const IR::IDeclaration *decl, cstring name) {
+void RenameMap::setNewName(const IR::IDeclaration *decl, cstring name, bool allowOverride) {
     CHECK_NULL(decl);
     BUG_CHECK(!name.isNullOrEmpty(), "Empty name");
     LOG1("Will rename " << dbp(decl) << " to " << name);
-    if (newName.find(decl) != newName.end()) BUG("%1%: already renamed", decl);
-    newName.emplace(decl, name);
+    BUG_CHECK(allowOverride || newName.find(decl) == newName.end(), "%1%: already renamed", decl);
+    newName.insert_or_assign(decl, name);
 }
 
 const IR::P4Action *RenameMap::actionCalled(const IR::MethodCallExpression *expression) const {
@@ -97,9 +97,9 @@ UniqueParameters::UniqueParameters(ReferenceMap *refMap, TypeMap *typeMap)
 
 IR::ID *RenameSymbols::getName() const {
     auto orig = getOriginal<IR::IDeclaration>();
-    if (!renameMap->toRename(orig)) return nullptr;
-    auto newName = renameMap->getName(orig);
-    auto name = new IR::ID(orig->getName().srcInfo, newName, orig->getName().originalName);
+    auto newName = renameMap->get(orig);
+    if (!newName.has_value()) return nullptr;
+    auto name = new IR::ID(orig->getName().srcInfo, *newName, orig->getName().originalName);
     return name;
 }
 
diff --git a/frontends/p4/uniqueNames.h b/frontends/p4/uniqueNames.h
index fa464e46255..3584cf3faeb 100644
--- a/frontends/p4/uniqueNames.h
+++ b/frontends/p4/uniqueNames.h
@@ -31,17 +31,33 @@ class RenameMap {
     std::map<const IR::MethodCallExpression *, const IR::P4Action *> actionCall;
 
  public:
-    void setNewName(const IR::IDeclaration *decl, cstring name);
+    /// @brief Add rename entry for the declaration to be named with the given name.
+    /// @param allowOverride If set to true, don't fail if a new name was already set but replace it
+    /// instead.
+    void setNewName(const IR::IDeclaration *decl, cstring name, bool allowOverride = false);
+
+    /// Get new name for the declaration, fails if none exists.
     cstring getName(const IR::IDeclaration *decl) const {
-        CHECK_NULL(decl);
-        BUG_CHECK(newName.find(decl) != newName.end(), "%1%: no new name", decl);
-        auto result = ::get(newName, decl);
-        return result;
+        auto n = get(decl);
+        BUG_CHECK(n.has_value(), "%1%: no new name", decl);
+        return *n;
     }
+
+    /// @returns true if there is a new name for the declaration, false otherwise.
     bool toRename(const IR::IDeclaration *decl) const {
         CHECK_NULL(decl);
         return newName.find(decl) != newName.end();
     }
+
+    /// Get new name for the declaration (wrapped in optional), or std::nullopt if there is none.
+    std::optional<cstring> get(const IR::IDeclaration *decl) const {
+        CHECK_NULL(decl);
+        if (auto it = newName.find(decl); it != newName.end()) {
+            return it->second;
+        }
+        return {};
+    }
+
     void foundInTable(const IR::P4Action *action);
     void markActionCall(const IR::P4Action *action, const IR::MethodCallExpression *call);
     const IR::P4Action *actionCalled(const IR::MethodCallExpression *expression) const;