[llvm] Add StringMap::lookup() overload to take a default value and return a reference to the value to avoid a copy (PR #115469)

[Boaz Brickner via llvm-commits]

https://github.com/bricknerb created https://github.com/llvm/llvm-project/pull/115469

Also use it in a few places in the code where it fits.

>From 468007d5bdcaf4308697dce682a4454c970272d9 Mon Sep 17 00:00:00 2001
From: Boaz Brickner <brickner at google.com>
Date: Fri, 8 Nov 2024 13:15:12 +0100
Subject: [PATCH] Add StringMap::lookup() overload to take a default value and
 return a reference to the value to avoid a copy Also use it in a few places
 in the code where it fits.

---
 llvm/include/llvm/ADT/StringMap.h         |  9 +++++++++
 llvm/lib/CodeGen/MachineOutliner.cpp      | 12 +++---------
 llvm/tools/llvm-rc/ResourceScriptStmt.cpp |  5 +----
 llvm/unittests/ADT/StringMapTest.cpp      |  5 +++++
 4 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/llvm/include/llvm/ADT/StringMap.h b/llvm/include/llvm/ADT/StringMap.h
index 9b58af73273913..be57cce851fa25 100644
--- a/llvm/include/llvm/ADT/StringMap.h
+++ b/llvm/include/llvm/ADT/StringMap.h
@@ -257,6 +257,15 @@ class LLVM_ALLOCATORHOLDER_EMPTYBASE StringMap
     return ValueTy();
   }
 
+  /// lookup - Return the entry for the specified key, or a default
+  /// provided value if no such entry exists.
+  const ValueTy &lookup(StringRef Key, const ValueTy& Value) const {
+    const_iterator Iter = find(Key);
+    if (Iter != end())
+      return Iter->second;
+    return Value;
+  }
+
   /// at - Return the entry for the specified key, or abort if no such
   /// entry exists.
   const ValueTy &at(StringRef Val) const {
diff --git a/llvm/lib/CodeGen/MachineOutliner.cpp b/llvm/lib/CodeGen/MachineOutliner.cpp
index 2ae6b4f2c64c9b..b89340d378c074 100644
--- a/llvm/lib/CodeGen/MachineOutliner.cpp
+++ b/llvm/lib/CodeGen/MachineOutliner.cpp
@@ -1289,15 +1289,9 @@ void MachineOutliner::emitInstrCountChangedRemark(
 
     std::string Fname = std::string(F.getName());
     unsigned FnCountAfter = MF->getInstructionCount();
-    unsigned FnCountBefore = 0;
-
-    // Check if the function was recorded before.
-    auto It = FunctionToInstrCount.find(Fname);
-
-    // Did we have a previously-recorded size? If yes, then set FnCountBefore
-    // to that.
-    if (It != FunctionToInstrCount.end())
-      FnCountBefore = It->second;
+    // Check if the function was recorded before and if yes, then set
+    // FnCountBefore to that.
+    unsigned FnCountBefore = FunctionToInstrCount.lookup(Fname, 0);
 
     // Compute the delta and emit a remark if there was a change.
     int64_t FnDelta = static_cast<int64_t>(FnCountAfter) -
diff --git a/llvm/tools/llvm-rc/ResourceScriptStmt.cpp b/llvm/tools/llvm-rc/ResourceScriptStmt.cpp
index a7f3df0863e74b..9cd645609a5f4d 100644
--- a/llvm/tools/llvm-rc/ResourceScriptStmt.cpp
+++ b/llvm/tools/llvm-rc/ResourceScriptStmt.cpp
@@ -228,10 +228,7 @@ const StringMap<VersionInfoFixedType> VersionInfoFixed::FixedFieldsInfoMap = {
 
 VersionInfoFixedType VersionInfoFixed::getFixedType(StringRef Type) {
   auto UpperType = Type.upper();
-  auto Iter = FixedFieldsInfoMap.find(UpperType);
-  if (Iter != FixedFieldsInfoMap.end())
-    return Iter->getValue();
-  return FtUnknown;
+  return FixedFieldsInfoMap.lookup(UpperType, FtUnknown);
 }
 
 bool VersionInfoFixed::isTypeSupported(VersionInfoFixedType Type) {
diff --git a/llvm/unittests/ADT/StringMapTest.cpp b/llvm/unittests/ADT/StringMapTest.cpp
index 35135cdb297e79..8959390c1c1a03 100644
--- a/llvm/unittests/ADT/StringMapTest.cpp
+++ b/llvm/unittests/ADT/StringMapTest.cpp
@@ -174,10 +174,15 @@ TEST_F(StringMapTest, SmallFullMapTest) {
 
   EXPECT_EQ(3u, Map.size());
   EXPECT_EQ(0, Map.lookup("eins"));
+  EXPECT_EQ(7, Map.lookup("eins", 7));
   EXPECT_EQ(2, Map.lookup("zwei"));
+  EXPECT_EQ(2, Map.lookup("zwei", 7));
   EXPECT_EQ(0, Map.lookup("drei"));
+  EXPECT_EQ(7, Map.lookup("drei", 7));
   EXPECT_EQ(4, Map.lookup("veir"));
+  EXPECT_EQ(4, Map.lookup("veir", 7));
   EXPECT_EQ(5, Map.lookup("funf"));
+  EXPECT_EQ(5, Map.lookup("funf", 7));
 }
 
 TEST_F(StringMapTest, CopyCtorTest) {