[llvm] [ReachingDefAnalysis][NFC] Fix management of MBBFrameObjsReachingDefs (PR #124943)

[Michael Maitland via llvm-commits]

https://github.com/michaelmaitland updated https://github.com/llvm/llvm-project/pull/124943

>From 07f09a4c2c6226a5a6bdd00edede86d255d78fc5 Mon Sep 17 00:00:00 2001
From: Michael Maitland <michaeltmaitland at gmail.com>
Date: Wed, 29 Jan 2025 08:39:09 -0800
Subject: [PATCH 1/2] [ReachingDefAnalysis][NFC] Fix management of
 MBBFrameObjsReachingDefs

---
 llvm/lib/CodeGen/ReachingDefAnalysis.cpp | 18 ++++++++++--------
 1 file changed, 10 insertions(+), 8 deletions(-)

diff --git a/llvm/lib/CodeGen/ReachingDefAnalysis.cpp b/llvm/lib/CodeGen/ReachingDefAnalysis.cpp
index 3d88c6815d51c9..269c8b8dbf53b0 100644
--- a/llvm/lib/CodeGen/ReachingDefAnalysis.cpp
+++ b/llvm/lib/CodeGen/ReachingDefAnalysis.cpp
@@ -147,17 +147,19 @@ void ReachingDefAnalysis::processDefs(MachineInstr *MI) {
       assert(FrameIndex >= 0 && "Can't handle negative frame indicies yet!");
       if (!isFIDef(*MI, FrameIndex, TII))
         continue;
+
+      int Key = FrameIndex - ObjectIndexBegin;
       if (MBBFrameObjsReachingDefs.contains(MBBNumber)) {
-        auto Frame2InstrIdx = MBBFrameObjsReachingDefs[MBBNumber];
-        if (Frame2InstrIdx.count(FrameIndex - ObjectIndexBegin) > 0)
-          Frame2InstrIdx[FrameIndex - ObjectIndexBegin].push_back(CurInstr);
+        auto &Frame2InstrIdx = MBBFrameObjsReachingDefs[MBBNumber];
+        if (Frame2InstrIdx.count(Key) > 0)
+          Frame2InstrIdx[Key].push_back(CurInstr);
         else
-          Frame2InstrIdx[FrameIndex - ObjectIndexBegin] = {CurInstr};
+          Frame2InstrIdx[Key] = {CurInstr};
       } else {
-        MBBFrameObjsReachingDefs[MBBNumber] = {
-            {FrameIndex - ObjectIndexBegin, {CurInstr}}};
+        MBBFrameObjsReachingDefs[MBBNumber] = {{Key, {CurInstr}}};
       }
     }
+
     if (!isValidRegDef(MO))
       continue;
     for (MCRegUnit Unit : TRI->regunits(MO.getReg().asMCReg())) {
@@ -348,8 +350,8 @@ int ReachingDefAnalysis::getReachingDef(MachineInstr *MI, Register Reg) const {
 
   if (Register::isStackSlot(Reg)) {
     int FrameIndex = Register::stackSlot2Index(Reg);
-    for (int Def : MBBFrameObjsReachingDefs.lookup(MBBNumber).lookup(
-             FrameIndex - ObjectIndexBegin)) {
+    int Key = FrameIndex - ObjectIndexBegin;
+    for (int Def : MBBFrameObjsReachingDefs.at(MBBNumber).at(Key)) {
       if (Def >= InstId)
         break;
       DefRes = Def;

>From 18aa1a136423225dcf4a4f849650e9a9fbfb7a23 Mon Sep 17 00:00:00 2001
From: Michael Maitland <michaeltmaitland at gmail.com>
Date: Wed, 29 Jan 2025 09:46:07 -0800
Subject: [PATCH 2/2] fixup! avoid assertion failure

---
 llvm/lib/CodeGen/ReachingDefAnalysis.cpp | 28 ++++++++++++++----------
 1 file changed, 17 insertions(+), 11 deletions(-)

diff --git a/llvm/lib/CodeGen/ReachingDefAnalysis.cpp b/llvm/lib/CodeGen/ReachingDefAnalysis.cpp
index 269c8b8dbf53b0..76560feb547c33 100644
--- a/llvm/lib/CodeGen/ReachingDefAnalysis.cpp
+++ b/llvm/lib/CodeGen/ReachingDefAnalysis.cpp
@@ -149,17 +149,12 @@ void ReachingDefAnalysis::processDefs(MachineInstr *MI) {
         continue;
 
       int Key = FrameIndex - ObjectIndexBegin;
-      if (MBBFrameObjsReachingDefs.contains(MBBNumber)) {
-        auto &Frame2InstrIdx = MBBFrameObjsReachingDefs[MBBNumber];
-        if (Frame2InstrIdx.count(Key) > 0)
-          Frame2InstrIdx[Key].push_back(CurInstr);
-        else
-          Frame2InstrIdx[Key] = {CurInstr};
-      } else {
-        MBBFrameObjsReachingDefs[MBBNumber] = {{Key, {CurInstr}}};
-      }
+      if (MBBFrameObjsReachingDefs.contains(MBBNumber) &&
+          MBBFrameObjsReachingDefs[MBBNumber].count(Key) > 0)
+        MBBFrameObjsReachingDefs[MBBNumber][Key].push_back(CurInstr);
+      else
+        MBBFrameObjsReachingDefs[MBBNumber][Key] = {CurInstr};
     }
-
     if (!isValidRegDef(MO))
       continue;
     for (MCRegUnit Unit : TRI->regunits(MO.getReg().asMCReg())) {
@@ -351,7 +346,18 @@ int ReachingDefAnalysis::getReachingDef(MachineInstr *MI, Register Reg) const {
   if (Register::isStackSlot(Reg)) {
     int FrameIndex = Register::stackSlot2Index(Reg);
     int Key = FrameIndex - ObjectIndexBegin;
-    for (int Def : MBBFrameObjsReachingDefs.at(MBBNumber).at(Key)) {
+
+    // Check that there was a reaching def.
+    auto Lookup1 = MBBFrameObjsReachingDefs.find(MBBNumber);
+    if (Lookup1 == MBBFrameObjsReachingDefs.end())
+      return LatestDef;
+    auto &InnerMap = Lookup1->second;
+    auto Lookup2 = InnerMap.find(Key);
+    if (Lookup2 == InnerMap.end())
+      return LatestDef;
+    auto &Defs = Lookup2->second;
+
+    for (int Def : Defs) {
       if (Def >= InstId)
         break;
       DefRes = Def;