[llvm] machinebackend pass test (PR #122487)

Kun Liu via llvm-commits llvm-commits at lists.llvm.org
Thu Jan 16 21:23:30 PST 2025


https://github.com/Ryan-hub-bit updated https://github.com/llvm/llvm-project/pull/122487

>From 9e95cde15ce30630e7118fe046b6dff2d7adf7d9 Mon Sep 17 00:00:00 2001
From: Liu <kliu14 at tulane.edu>
Date: Fri, 10 Jan 2025 10:21:22 -0600
Subject: [PATCH 1/7] machinebackend pass test

---
 llvm/lib/Target/X86/X86MatchJumptablePass.cpp | 44 +++++++++++++++++++
 llvm/lib/Target/X86/X86MatchJumptablePass.h   |  0
 2 files changed, 44 insertions(+)
 create mode 100644 llvm/lib/Target/X86/X86MatchJumptablePass.cpp
 create mode 100644 llvm/lib/Target/X86/X86MatchJumptablePass.h

diff --git a/llvm/lib/Target/X86/X86MatchJumptablePass.cpp b/llvm/lib/Target/X86/X86MatchJumptablePass.cpp
new file mode 100644
index 00000000000000..bfb6b1b7f89f18
--- /dev/null
+++ b/llvm/lib/Target/X86/X86MatchJumptablePass.cpp
@@ -0,0 +1,44 @@
+#include "X86.h"
+#include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/MachineInstr.h"
+#include "llvm/CodeGen/MachineRegisterInfo.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/raw_ostream.h"
+
+
+using namespace llvm;
+
+namespace {
+  class X86MatchJumptablePass : public MachineFunctionPass {
+  public:
+    static char ID;
+
+    X86MatchJumptablePass() : MachineFunctionPass(ID) {}
+
+    bool runOnMachineFunction(MachineFunction &MF) override {
+      LLVM_DEBUG(dbgs() << "Running X86MyBackendPass on function: "
+                        << MF.getName() << "\n");
+
+      // Example: Iterate through instructions
+      for (auto &MBB : MF) {
+        for (auto &MI : MBB) {
+          // Process instructions here
+          LLVM_DEBUG(dbgs() << "Instruction: " << MI << "\n");
+        }
+      }
+
+      return false; // Return true if the pass modifies the function
+    }
+
+    StringRef getPassName() const override {
+      return "X86 My Backend Pass";
+    }
+  };
+}
+
+char X86MatchJumptablePass::ID = 0;
+
+// Register the pass
+FunctionPass *llvm::createX86MatchJumptablePass() {
+  return new X86MatchJumptablePass();
+}
diff --git a/llvm/lib/Target/X86/X86MatchJumptablePass.h b/llvm/lib/Target/X86/X86MatchJumptablePass.h
new file mode 100644
index 00000000000000..e69de29bb2d1d6

>From 27bee8dbc8fa03e331302e799429b7450c2b1166 Mon Sep 17 00:00:00 2001
From: ryan <liu864579887 at gmail.com>
Date: Fri, 10 Jan 2025 17:01:39 -0600
Subject: [PATCH 2/7] good

---
 .../llvm/Transforms/IPO/JumpTableFinder.h     | 62 +++++++++++++++++++
 llvm/lib/Target/X86/CMakeLists.txt            |  1 +
 llvm/lib/Target/X86/X86MatchJumptablePass.cpp | 15 +++--
 llvm/lib/Target/X86/X86MatchJumptablePass.h   | 14 +++++
 llvm/lib/Target/X86/X86TargetMachine.cpp      |  3 +-
 5 files changed, 90 insertions(+), 5 deletions(-)
 create mode 100644 llvm/include/llvm/Transforms/IPO/JumpTableFinder.h

diff --git a/llvm/include/llvm/Transforms/IPO/JumpTableFinder.h b/llvm/include/llvm/Transforms/IPO/JumpTableFinder.h
new file mode 100644
index 00000000000000..9d2fb3b405f8c8
--- /dev/null
+++ b/llvm/include/llvm/Transforms/IPO/JumpTableFinder.h
@@ -0,0 +1,62 @@
+// #ifndef LLVM_ANALYSIS_JUMPTABLEFINDERPASS_H
+// #define LLVM_ANALYSIS_JUMPTABLEFINDERPASS_H
+
+// #include "llvm/IR/PassManager.h"
+// #include "llvm/IR/Module.h"
+// #include "llvm/Support/raw_ostream.h"
+// #include <set>
+
+// namespace llvm {
+
+// class JumptableFinderPass : public PassInfoMixin<JumptableFinderPass> {
+// public:
+//     /// Main entry point for the pass. Analyzes the module to find and analyze jump tables.
+//     PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM);
+//     /// Implementation of the jump table finder.
+//     void jumptableFinderImpl(Module &M);
+
+//     /// Analyze a SwitchInst for potential jump table patterns.
+//     void findJumpTableFromSwitch(SwitchInst *SI);
+
+//     /// Analyze a GetElementPtrInst for jump table patterns.
+//     void analyzeJumpTable(GetElementPtrInst *GEP);
+
+//     /// Analyze the index computation of a jump table.
+//     void analyzeIndex(Value *Index);
+
+//     /// Find all potential targets for a jump table.
+//     void findTargets(GetElementPtrInst *GEP, std::set<BasicBlock*> &Targets);
+
+//     /// Check the density of a SwitchInst's cases to determine if it forms a jump table.
+//     bool checkDensity(SwitchInst *SI);
+
+//     /// Check if a GetElementPtrInst leads to an indirect branch.
+//     bool leadsToIndirectBranch(GetElementPtrInst *GEP);
+// };
+
+// } // namespace llvm
+
+// #endif // LLVM_ANALYSIS_JUMPTABLEFINDERPASS_H
+
+#ifndef LLVM_TRANSFORMS_IPO_JUMPTABLEFINDER_H
+#define LLVM_TRANSFORMS_IPO_JUMPTABLEFINDER_H
+
+#include "llvm/IR/Module.h"
+#include "llvm/IR/PassManager.h" // For PassInfoMixin and PreservedAnalyses
+#include "llvm/IR/Function.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/Support/raw_ostream.h"
+#include <set>
+
+namespace llvm {
+
+class JumptableFinderPass : public PassInfoMixin<JumptableFinderPass> {
+public:
+    // Entry point for the pass
+    PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
+};
+
+} // namespace llvm
+
+#endif // LLVM_TRANSFORMS_IPO_JUMPTABLEFINDER_H
diff --git a/llvm/lib/Target/X86/CMakeLists.txt b/llvm/lib/Target/X86/CMakeLists.txt
index 9553a8619feb51..4be4576a6a9057 100644
--- a/llvm/lib/Target/X86/CMakeLists.txt
+++ b/llvm/lib/Target/X86/CMakeLists.txt
@@ -23,6 +23,7 @@ tablegen(LLVM X86GenFoldTables.inc -gen-x86-fold-tables -asmwriternum=1)
 add_public_tablegen_target(X86CommonTableGen)
 
 set(sources
+X86MatchJumptablePass.cpp
   X86ArgumentStackSlotRebase.cpp
   X86AsmPrinter.cpp
   X86AvoidTrailingCall.cpp
diff --git a/llvm/lib/Target/X86/X86MatchJumptablePass.cpp b/llvm/lib/Target/X86/X86MatchJumptablePass.cpp
index bfb6b1b7f89f18..c3683dacbdf8fe 100644
--- a/llvm/lib/Target/X86/X86MatchJumptablePass.cpp
+++ b/llvm/lib/Target/X86/X86MatchJumptablePass.cpp
@@ -4,7 +4,9 @@
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
+#include "X86MatchJumptablePass.h"
 
+#define DEBUG_TYPE "x86-my-pass"
 
 using namespace llvm;
 
@@ -38,7 +40,12 @@ namespace {
 
 char X86MatchJumptablePass::ID = 0;
 
-// Register the pass
-FunctionPass *llvm::createX86MatchJumptablePass() {
-  return new X86MatchJumptablePass();
-}
+// Ensure the function is in the llvm namespace
+namespace llvm {
+  
+  // Define the pass
+  FunctionPass *createX86MatchJumptablePass() {
+    return new X86MatchJumptablePass();
+  }
+
+} // end llvm namespace
diff --git a/llvm/lib/Target/X86/X86MatchJumptablePass.h b/llvm/lib/Target/X86/X86MatchJumptablePass.h
index e69de29bb2d1d6..67233bf729aec6 100644
--- a/llvm/lib/Target/X86/X86MatchJumptablePass.h
+++ b/llvm/lib/Target/X86/X86MatchJumptablePass.h
@@ -0,0 +1,14 @@
+#ifndef LLVM_LIB_TARGET_X86_X86MATCHJUMPTABLEPASS_H
+#define LLVM_LIB_TARGET_X86_X86MATCHJUMPTABLEPASS_H
+
+
+#pragma once
+#include "llvm/CodeGen/MachineFunctionPass.h"
+
+namespace llvm {
+
+  FunctionPass *createX86MatchJumptablePass();
+
+} // end namespace llvm
+
+#endif // LLVM_LIB_TARGET_X86_X86MATCHJUMPTABLEPASS_H
\ No newline at end of file
diff --git a/llvm/lib/Target/X86/X86TargetMachine.cpp b/llvm/lib/Target/X86/X86TargetMachine.cpp
index 20dfdd27b33df6..b9c13fda07dfc7 100644
--- a/llvm/lib/Target/X86/X86TargetMachine.cpp
+++ b/llvm/lib/Target/X86/X86TargetMachine.cpp
@@ -12,6 +12,7 @@
 
 #include "X86TargetMachine.h"
 #include "MCTargetDesc/X86MCTargetDesc.h"
+#include "X86MatchJumptablePass.h"  // Include the header with the function declaration
 #include "TargetInfo/X86TargetInfo.h"
 #include "X86.h"
 #include "X86MachineFunctionInfo.h"
@@ -464,7 +465,7 @@ MachineFunctionInfo *X86TargetMachine::createMachineFunctionInfo(
 
 void X86PassConfig::addIRPasses() {
   addPass(createAtomicExpandLegacyPass());
-
+  addPass(llvm::createX86MatchJumptablePass());
   // We add both pass anyway and when these two passes run, we skip the pass
   // based on the option level and option attribute.
   addPass(createX86LowerAMXIntrinsicsPass());

>From 0dc2ceac51b6c4182a8816b215146b3b7a68c1bd Mon Sep 17 00:00:00 2001
From: Liu <kliu14 at tulane.edu>
Date: Fri, 10 Jan 2025 20:05:42 -0600
Subject: [PATCH 3/7] jumptable

---
 llvm/lib/Target/X86/CMakeLists.txt            |   2 +-
 llvm/lib/Target/X86/X86MatchJumptablePass.cpp | 102 +++++++++++++++---
 llvm/lib/Target/X86/X86MatchJumptablePass.h   |  10 +-
 llvm/lib/Target/X86/X86TargetMachine.cpp      |   5 +-
 4 files changed, 99 insertions(+), 20 deletions(-)

diff --git a/llvm/lib/Target/X86/CMakeLists.txt b/llvm/lib/Target/X86/CMakeLists.txt
index 4be4576a6a9057..f5dcf13e47ed25 100644
--- a/llvm/lib/Target/X86/CMakeLists.txt
+++ b/llvm/lib/Target/X86/CMakeLists.txt
@@ -23,7 +23,7 @@ tablegen(LLVM X86GenFoldTables.inc -gen-x86-fold-tables -asmwriternum=1)
 add_public_tablegen_target(X86CommonTableGen)
 
 set(sources
-X86MatchJumptablePass.cpp
+  X86MatchJumptablePass.cpp
   X86ArgumentStackSlotRebase.cpp
   X86AsmPrinter.cpp
   X86AvoidTrailingCall.cpp
diff --git a/llvm/lib/Target/X86/X86MatchJumptablePass.cpp b/llvm/lib/Target/X86/X86MatchJumptablePass.cpp
index c3683dacbdf8fe..d8263af59a855d 100644
--- a/llvm/lib/Target/X86/X86MatchJumptablePass.cpp
+++ b/llvm/lib/Target/X86/X86MatchJumptablePass.cpp
@@ -6,7 +6,7 @@
 #include "llvm/Support/raw_ostream.h"
 #include "X86MatchJumptablePass.h"
 
-#define DEBUG_TYPE "x86-my-pass"
+#define DEBUG_TYPE "match-jump-table"
 
 using namespace llvm;
 
@@ -18,23 +18,97 @@ namespace {
     X86MatchJumptablePass() : MachineFunctionPass(ID) {}
 
     bool runOnMachineFunction(MachineFunction &MF) override {
-      LLVM_DEBUG(dbgs() << "Running X86MyBackendPass on function: "
-                        << MF.getName() << "\n");
-
-      // Example: Iterate through instructions
-      for (auto &MBB : MF) {
-        for (auto &MI : MBB) {
-          // Process instructions here
-          LLVM_DEBUG(dbgs() << "Instruction: " << MI << "\n");
+    LLVM_DEBUG(dbgs() << "Analyzing jump tables in function: " << MF.getName() << "\n");
+
+    // Get jump table information
+    MachineJumpTableInfo *JumpTableInfo = MF.getJumpTableInfo();
+    if (!JumpTableInfo) {
+      LLVM_DEBUG(dbgs() << "No jump tables in this function.\n");
+      return false;
+    }
+    // Assuming JumpTableInfo is available
+    for (unsigned JTIndex = 0; JTIndex < JumpTableInfo->getJumpTables().size(); ++JTIndex) {
+  const MachineJumpTableEntry &JTEntry = JumpTableInfo->getJumpTables()[JTIndex];
+  
+  LLVM_DEBUG(dbgs() << "Jump Table #" << JTIndex << " Base Address: " << JTEntry.BaseAddress << "\n");
+
+  // Iterate through the entries (target basic blocks) in this jump table
+  for (auto *MBB : JTEntry.MBBs) {
+    LLVM_DEBUG(dbgs() << "  Target BasicBlock: " << MBB->getName() << " Address: " << MBB->getAddress() << "\n");
+     }
+    
+
+      // Trace potential indirect jumps related to this jump table
+      traceIndirectJumps(MF, JTIndex, JumpTableInfo);
+    }
+    return false;
+
+}
+
+  void traceIndirectJumps(MachineFunction &MF, unsigned JTIndex, MachineJumpTableInfo *JumpTableInfo) {
+    const MachineJumpTableEntry &JTEntry = JumpTableInfo->getJumpTables()[JTIndex];
+
+    for (auto &MBB : MF) {
+      for (auto &MI : MBB) {
+        if (MI.isIndirectBranch()) {
+          LLVM_DEBUG(dbgs() << "Found indirect jump: " << MI << "\n");
+
+          // Analyze data flow to check if this jump is related to the jump table
+          if (isJumpTableRelated(MI, JTEntry, MF)) {
+            LLVM_DEBUG(dbgs() << "This indirect jump is related to Jump Table #" << JTIndex << "\n");
+          }
         }
       }
+    }
+  }
+
+  bool isJumpTableRelated(MachineInstr &MI, const MachineJumpTableEntry &JTEntry, MachineFunction &MF) {
+  for (unsigned OpIdx = 0; OpIdx < MI.getNumOperands(); ++OpIdx) {
+    const MachineOperand &Op = MI.getOperand(OpIdx);
 
-      return false; // Return true if the pass modifies the function
+    if (Op.isReg()) {
+      Register Reg = Op.getReg();
+      MachineRegisterInfo &MRI = MF.getRegInfo();
+      // Check if any of the definitions of the register are related to a jump table load
+      for (MachineInstr &DefMI : MRI.def_instructions(Reg)) {
+        if (isJumpTableLoad(DefMI, JTEntry)) {
+          return true;
+        }
+      }
+    } else if (Op.isImm()) {
+      // Check if the immediate operand might be an offset/index into the jump table
+      int64_t ImmValue = Op.getImm();
+      
+      // For example, if the jump table has 10 entries, check if the immediate is between 0 and 9
+      if (ImmValue >= 0 && ImmValue < JTEntry.MBBs.size()) {
+        // This immediate value could be an index into the jump table
+        LLVM_DEBUG(dbgs() << "Immediate operand is a possible jump table index: " << ImmValue << "\n");
+        return true;
+      }
     }
+  }
+  return false;
+}
 
-    StringRef getPassName() const override {
-      return "X86 My Backend Pass";
+bool isJumpTableLoad(MachineInstr &MI, const MachineJumpTableEntry &JTEntry) {
+    if (MI.mayLoad()) {
+        for (unsigned i = 0; i < MI.getNumOperands(); ++i) {
+            const MachineOperand &Op = MI.getOperand(i);
+            if (Op.isGlobal() && Op.getGlobal() == JTEntry.BaseAddress) {
+                return true;
+            }
+        }
     }
+    return false;
+}
+
+  StringRef getPassName() const override {
+    return "Match Jump Table Pass";
+  }
+
+    // StringRef getPassName() const override {
+    //   return "X86 My Backend Pass";
+    // }
   };
 }
 
@@ -49,3 +123,7 @@ namespace llvm {
   }
 
 } // end llvm namespace
+
+static RegisterPass<X86MatchJumptablePass> X("match-jump-table", "Match Jump Table Pass", false, false);
+
+
diff --git a/llvm/lib/Target/X86/X86MatchJumptablePass.h b/llvm/lib/Target/X86/X86MatchJumptablePass.h
index 67233bf729aec6..ac6fc315b025ad 100644
--- a/llvm/lib/Target/X86/X86MatchJumptablePass.h
+++ b/llvm/lib/Target/X86/X86MatchJumptablePass.h
@@ -1,14 +1,14 @@
 #ifndef LLVM_LIB_TARGET_X86_X86MATCHJUMPTABLEPASS_H
 #define LLVM_LIB_TARGET_X86_X86MATCHJUMPTABLEPASS_H
 
-
-#pragma once
 #include "llvm/CodeGen/MachineFunctionPass.h"
 
 namespace llvm {
 
-  FunctionPass *createX86MatchJumptablePass();
+/// \brief Creates the X86MatchJumptablePass.
+/// This pass analyzes and processes jump tables in X86 backend code generation.
+FunctionPass *createX86MatchJumptablePass();
 
-} // end namespace llvm
+} // namespace llvm
 
-#endif // LLVM_LIB_TARGET_X86_X86MATCHJUMPTABLEPASS_H
\ No newline at end of file
+#endif // LLVM_LIB_TARGET_X86_X86MATCHJUMPTABLEPASS_H
diff --git a/llvm/lib/Target/X86/X86TargetMachine.cpp b/llvm/lib/Target/X86/X86TargetMachine.cpp
index b9c13fda07dfc7..7faeeb2af64e8e 100644
--- a/llvm/lib/Target/X86/X86TargetMachine.cpp
+++ b/llvm/lib/Target/X86/X86TargetMachine.cpp
@@ -11,8 +11,8 @@
 //===----------------------------------------------------------------------===//
 
 #include "X86TargetMachine.h"
+#include "X86MatchJumptablePass.h"
 #include "MCTargetDesc/X86MCTargetDesc.h"
-#include "X86MatchJumptablePass.h"  // Include the header with the function declaration
 #include "TargetInfo/X86TargetInfo.h"
 #include "X86.h"
 #include "X86MachineFunctionInfo.h"
@@ -465,7 +465,7 @@ MachineFunctionInfo *X86TargetMachine::createMachineFunctionInfo(
 
 void X86PassConfig::addIRPasses() {
   addPass(createAtomicExpandLegacyPass());
-  addPass(llvm::createX86MatchJumptablePass());
+
   // We add both pass anyway and when these two passes run, we skip the pass
   // based on the option level and option attribute.
   addPass(createX86LowerAMXIntrinsicsPass());
@@ -597,6 +597,7 @@ void X86PassConfig::addPreEmitPass() {
     addPass(createBreakFalseDeps());
   }
 
+  addPass(createX86MatchJumptablePass());
   addPass(createX86IndirectBranchTrackingPass());
 
   addPass(createX86IssueVZeroUpperPass());

>From 65c7421be86c02e11d1e0e6d51b556a22063ba94 Mon Sep 17 00:00:00 2001
From: ryan <liu864579887 at gmail.com>
Date: Fri, 10 Jan 2025 21:50:10 -0600
Subject: [PATCH 4/7] test done

---
 llvm/lib/Target/X86/X86MatchJumptablePass.cpp | 307 +++++++++++++-----
 1 file changed, 234 insertions(+), 73 deletions(-)

diff --git a/llvm/lib/Target/X86/X86MatchJumptablePass.cpp b/llvm/lib/Target/X86/X86MatchJumptablePass.cpp
index d8263af59a855d..c43d7c5b676b49 100644
--- a/llvm/lib/Target/X86/X86MatchJumptablePass.cpp
+++ b/llvm/lib/Target/X86/X86MatchJumptablePass.cpp
@@ -1,23 +1,153 @@
+// #include "X86.h"
+// #include "llvm/CodeGen/MachineFunctionPass.h"
+// #include "llvm/CodeGen/MachineInstr.h"
+// #include "llvm/CodeGen/MachineRegisterInfo.h"
+// #include "llvm/Support/Debug.h"
+// #include "llvm/Support/raw_ostream.h"
+// #include "X86MatchJumptablePass.h"
+// #include "llvm/CodeGen/MachineJumpTableInfo.h"
+
+// #define DEBUG_TYPE "match-jump-table"
+
+// using namespace llvm;
+
+// namespace {
+//   class X86MatchJumptablePass : public MachineFunctionPass {
+//   public:
+//     static char ID;
+
+//     X86MatchJumptablePass() : MachineFunctionPass(ID) {}
+
+//     bool runOnMachineFunction(MachineFunction &MF) override {
+//     LLVM_DEBUG(dbgs() << "Analyzing jump tables in function: " << MF.getName() << "\n");
+
+//     // Get jump table information
+//     MachineJumpTableInfo *JumpTableInfo = MF.getJumpTableInfo();
+//     if (!JumpTableInfo) {
+//       LLVM_DEBUG(dbgs() << "No jump tables in this function.\n");
+//       return false;
+//     }
+//     // Assuming JumpTableInfo is available
+//     for (unsigned JTIndex = 0; JTIndex < JumpTableInfo->getJumpTables().size(); ++JTIndex) {
+//   const MachineJumpTableEntry &JTEntry = JumpTableInfo->getJumpTables()[JTIndex];
+  
+//   LLVM_DEBUG(dbgs() << "Jump Table #" << JTIndex << " Base Address: " << JTEntry.BaseAddress << "\n");
+
+//   // Iterate through the entries (target basic blocks) in this jump table
+//   for (auto *MBB : JTEntry.MBBs) {
+//     LLVM_DEBUG(dbgs() << "  Target BasicBlock: " << MBB->getName() << " Address: " << MBB->getAddress() << "\n");
+//      }
+    
+
+//       // Trace potential indirect jumps related to this jump table
+//       traceIndirectJumps(MF, JTIndex, JumpTableInfo);
+//     }
+//     return false;
+
+// }
+
+//   void traceIndirectJumps(MachineFunction &MF, unsigned JTIndex, MachineJumpTableInfo *JumpTableInfo) {
+//     const MachineJumpTableEntry &JTEntry = JumpTableInfo->getJumpTables()[JTIndex];
+
+//     for (auto &MBB : MF) {
+//       for (auto &MI : MBB) {
+//         if (MI.isIndirectBranch()) {
+//           LLVM_DEBUG(dbgs() << "Found indirect jump: " << MI << "\n");
+
+//           // Analyze data flow to check if this jump is related to the jump table
+//           if (isJumpTableRelated(MI, JTEntry, MF)) {
+//             LLVM_DEBUG(dbgs() << "This indirect jump is related to Jump Table #" << JTIndex << "\n");
+//           }
+//         }
+//       }
+//     }
+//   }
+
+//   bool isJumpTableRelated(MachineInstr &MI, const MachineJumpTableEntry &JTEntry, MachineFunction &MF) {
+//   for (unsigned OpIdx = 0; OpIdx < MI.getNumOperands(); ++OpIdx) {
+//     const MachineOperand &Op = MI.getOperand(OpIdx);
+
+//     if (Op.isReg()) {
+//       Register Reg = Op.getReg();
+//       MachineRegisterInfo &MRI = MF.getRegInfo();
+//       // Check if any of the definitions of the register are related to a jump table load
+//       for (MachineInstr &DefMI : MRI.def_instructions(Reg)) {
+//         if (isJumpTableLoad(DefMI, JTEntry)) {
+//           return true;
+//         }
+//       }
+//     } else if (Op.isImm()) {
+//       // Check if the immediate operand might be an offset/index into the jump table
+//       int64_t ImmValue = Op.getImm();
+      
+//       // For example, if the jump table has 10 entries, check if the immediate is between 0 and 9
+//       if (ImmValue >= 0 && ImmValue < JTEntry.MBBs.size()) {
+//         // This immediate value could be an index into the jump table
+//         LLVM_DEBUG(dbgs() << "Immediate operand is a possible jump table index: " << ImmValue << "\n");
+//         return true;
+//       }
+//     }
+//   }
+//   return false;
+// }
+
+// bool isJumpTableLoad(MachineInstr &MI, const MachineJumpTableEntry &JTEntry) {
+//     if (MI.mayLoad()) {
+//         for (unsigned i = 0; i < MI.getNumOperands(); ++i) {
+//             const MachineOperand &Op = MI.getOperand(i);
+//             if (Op.isGlobal() && Op.getGlobal() == JTEntry.BaseAddress) {
+//                 return true;
+//             }
+//         }
+//     }
+//     return false;
+// }
+
+//   StringRef getPassName() const override {
+//     return "Match Jump Table Pass";
+//   }
+
+//     // StringRef getPassName() const override {
+//     //   return "X86 My Backend Pass";
+//     // }
+//   };
+// }
+
+// char X86MatchJumptablePass::ID = 0;
+
+// // Ensure the function is in the llvm namespace
+// namespace llvm {
+  
+//   // Define the pass
+//   FunctionPass *createX86MatchJumptablePass() {
+//     return new X86MatchJumptablePass();
+//   }
+
+// } // end llvm namespace
+
+// static RegisterPass<X86MatchJumptablePass> X("match-jump-table", "Match Jump Table Pass", false, false);
+
+
 #include "X86.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineInstr.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
-#include "X86MatchJumptablePass.h"
+#include "llvm/CodeGen/MachineJumpTableInfo.h"
 
 #define DEBUG_TYPE "match-jump-table"
 
 using namespace llvm;
 
 namespace {
-  class X86MatchJumptablePass : public MachineFunctionPass {
-  public:
-    static char ID;
+class X86MatchJumptablePass : public MachineFunctionPass {
+public:
+  static char ID;
 
-    X86MatchJumptablePass() : MachineFunctionPass(ID) {}
+  X86MatchJumptablePass() : MachineFunctionPass(ID) {}
 
-    bool runOnMachineFunction(MachineFunction &MF) override {
+  bool runOnMachineFunction(MachineFunction &MF) override {
     LLVM_DEBUG(dbgs() << "Analyzing jump tables in function: " << MF.getName() << "\n");
 
     // Get jump table information
@@ -26,104 +156,135 @@ namespace {
       LLVM_DEBUG(dbgs() << "No jump tables in this function.\n");
       return false;
     }
-    // Assuming JumpTableInfo is available
+
     for (unsigned JTIndex = 0; JTIndex < JumpTableInfo->getJumpTables().size(); ++JTIndex) {
-  const MachineJumpTableEntry &JTEntry = JumpTableInfo->getJumpTables()[JTIndex];
-  
-  LLVM_DEBUG(dbgs() << "Jump Table #" << JTIndex << " Base Address: " << JTEntry.BaseAddress << "\n");
+      const MachineJumpTableEntry &JTEntry = JumpTableInfo->getJumpTables()[JTIndex];
 
-  // Iterate through the entries (target basic blocks) in this jump table
-  for (auto *MBB : JTEntry.MBBs) {
-    LLVM_DEBUG(dbgs() << "  Target BasicBlock: " << MBB->getName() << " Address: " << MBB->getAddress() << "\n");
-     }
-    
+      LLVM_DEBUG(dbgs() << "Jump Table #" << JTIndex << " contains " << JTEntry.MBBs.size()
+                        << " entries.\n");
+
+      // Iterate through the entries (target basic blocks) in this jump table
+      for (auto *MBB : JTEntry.MBBs) {
+        if (MBB) {
+          LLVM_DEBUG(dbgs() << "  Target BasicBlock: " << MBB->getName() << "\n");
+        }
+      }
 
-      // Trace potential indirect jumps related to this jump table
-      traceIndirectJumps(MF, JTIndex, JumpTableInfo);
+      // Assuming you have access to MF, JTIndex, and JumpTableInfo
+    MachineInstr* indirectJumpInstr = traceIndirectJumps(MF, JTIndex, JumpTableInfo);
+
+      if (indirectJumpInstr) {
+          // Handle the found indirect jump instruction
+          dbgs() << "Indirect jump found at address: " << indirectJumpInstr << "\n";
+          for (auto &MBB : JTEntry.MBBs) {
+          LLVM_DEBUG(dbgs() << "Address of MBB: " << &MBB << "\n"); // Print the address of the MBB
+          // Optionally print the name and instructions inside the MBB
+         // LLVM_DEBUG(dbgs() << "MBB: " << MBB.getName() << "\n");
+          // for (auto &MI : MBB) {
+          //     LLVM_DEBUG(dbgs() << "  Instruction: " << MI << "\n");
+          // }
+        }
+      } 
     }
+
     return false;
+  }
 
-}
+  // void traceIndirectJumps(MachineFunction &MF, unsigned JTIndex, MachineJumpTableInfo *JumpTableInfo) {
+  //   const MachineJumpTableEntry &JTEntry = JumpTableInfo->getJumpTables()[JTIndex];
+
+  //   for (auto &MBB : MF) {
+  //     for (auto &MI : MBB) {
+  //       if (MI.isIndirectBranch()) {
+  //         LLVM_DEBUG(dbgs() << "Found indirect jump: " << MI << "\n");
 
-  void traceIndirectJumps(MachineFunction &MF, unsigned JTIndex, MachineJumpTableInfo *JumpTableInfo) {
+  //         // Analyze data flow to check if this jump is related to the jump table
+  //         if (isJumpTableRelated(MI, JTEntry, MF)) {
+  //           LLVM_DEBUG(dbgs() << "This indirect jump is related to Jump Table #" << JTIndex << "\n");
+  //         }
+  //       }
+  //     }
+  //   }
+  // }
+  MachineInstr* traceIndirectJumps(MachineFunction &MF, unsigned JTIndex, MachineJumpTableInfo *JumpTableInfo) {
     const MachineJumpTableEntry &JTEntry = JumpTableInfo->getJumpTables()[JTIndex];
 
     for (auto &MBB : MF) {
-      for (auto &MI : MBB) {
-        if (MI.isIndirectBranch()) {
-          LLVM_DEBUG(dbgs() << "Found indirect jump: " << MI << "\n");
+        for (auto &MI : MBB) {
+            if (MI.isIndirectBranch()) {
+                LLVM_DEBUG(dbgs() << "Found indirect jump: " << MI << "\n");
 
-          // Analyze data flow to check if this jump is related to the jump table
-          if (isJumpTableRelated(MI, JTEntry, MF)) {
-            LLVM_DEBUG(dbgs() << "This indirect jump is related to Jump Table #" << JTIndex << "\n");
-          }
+                // Analyze data flow to check if this jump is related to the jump table
+                if (isJumpTableRelated(MI, JTEntry, MF)) {
+                    LLVM_DEBUG(dbgs() << "This indirect jump is related to Jump Table #" << JTIndex << "\n");
+                    
+                    // Return the address of the indirect jump (MI)
+                    return &MI;
+                }
+            }
         }
-      }
     }
-  }
+
+    // Return nullptr if no indirect jump is found
+    return nullptr;
+}
 
   bool isJumpTableRelated(MachineInstr &MI, const MachineJumpTableEntry &JTEntry, MachineFunction &MF) {
-  for (unsigned OpIdx = 0; OpIdx < MI.getNumOperands(); ++OpIdx) {
-    const MachineOperand &Op = MI.getOperand(OpIdx);
-
-    if (Op.isReg()) {
-      Register Reg = Op.getReg();
-      MachineRegisterInfo &MRI = MF.getRegInfo();
-      // Check if any of the definitions of the register are related to a jump table load
-      for (MachineInstr &DefMI : MRI.def_instructions(Reg)) {
-        if (isJumpTableLoad(DefMI, JTEntry)) {
+    for (unsigned OpIdx = 0; OpIdx < MI.getNumOperands(); ++OpIdx) {
+      const MachineOperand &Op = MI.getOperand(OpIdx);
+
+      if (Op.isReg()) {
+        Register Reg = Op.getReg();
+        MachineRegisterInfo &MRI = MF.getRegInfo();
+        // Check if any of the definitions of the register are related to a jump table load
+        for (MachineInstr &DefMI : MRI.def_instructions(Reg)) {
+          if (isJumpTableLoad(DefMI, JTEntry)) {
+            return true;
+          }
+        }
+      } else if (Op.isImm()) {
+        // Check if the immediate operand might be an offset/index into the jump table
+        int64_t ImmValue = Op.getImm();
+        if (ImmValue >= 0 && ImmValue < JTEntry.MBBs.size()) {
+          // This immediate value could be an index into the jump table
+          LLVM_DEBUG(dbgs() << "Immediate operand is a possible jump table index: " << ImmValue
+                            << "\n");
           return true;
         }
       }
-    } else if (Op.isImm()) {
-      // Check if the immediate operand might be an offset/index into the jump table
-      int64_t ImmValue = Op.getImm();
-      
-      // For example, if the jump table has 10 entries, check if the immediate is between 0 and 9
-      if (ImmValue >= 0 && ImmValue < JTEntry.MBBs.size()) {
-        // This immediate value could be an index into the jump table
-        LLVM_DEBUG(dbgs() << "Immediate operand is a possible jump table index: " << ImmValue << "\n");
-        return true;
-      }
     }
+    return false;
   }
-  return false;
-}
 
-bool isJumpTableLoad(MachineInstr &MI, const MachineJumpTableEntry &JTEntry) {
+  bool isJumpTableLoad(MachineInstr &MI, const MachineJumpTableEntry &JTEntry) {
     if (MI.mayLoad()) {
-        for (unsigned i = 0; i < MI.getNumOperands(); ++i) {
-            const MachineOperand &Op = MI.getOperand(i);
-            if (Op.isGlobal() && Op.getGlobal() == JTEntry.BaseAddress) {
-                return true;
-            }
+      for (unsigned i = 0; i < MI.getNumOperands(); ++i) {
+        const MachineOperand &Op = MI.getOperand(i);
+          if (Op.getType() == MachineOperand::MO_JumpTableIndex) {
+          unsigned JTIndex = Op.getIndex();
+          if (JTIndex < JTEntry.MBBs.size()) {
+            LLVM_DEBUG(dbgs() << "Instruction loads from Jump Table index: " << JTIndex << "\n");
+            return true;
+          }
         }
+      }
     }
     return false;
-}
-
-  StringRef getPassName() const override {
-    return "Match Jump Table Pass";
   }
 
-    // StringRef getPassName() const override {
-    //   return "X86 My Backend Pass";
-    // }
-  };
-}
+  StringRef getPassName() const override { return "Match Jump Table Pass"; }
+};
+
+} // namespace
 
 char X86MatchJumptablePass::ID = 0;
 
-// Ensure the function is in the llvm namespace
 namespace llvm {
-  
-  // Define the pass
-  FunctionPass *createX86MatchJumptablePass() {
-    return new X86MatchJumptablePass();
-  }
-
-} // end llvm namespace
 
-static RegisterPass<X86MatchJumptablePass> X("match-jump-table", "Match Jump Table Pass", false, false);
+// Define the pass
+FunctionPass *createX86MatchJumptablePass() { return new X86MatchJumptablePass(); }
 
+} // namespace llvm
 
+static RegisterPass<X86MatchJumptablePass> X("match-jump-table", "Match Jump Table Pass", false,
+                                             false);

>From 9dd981b98af2b53d024ea7a10064b49b8daa99f9 Mon Sep 17 00:00:00 2001
From: ryan <liu864579887 at gmail.com>
Date: Sun, 12 Jan 2025 18:40:37 -0600
Subject: [PATCH 5/7] clang jump table pattern test successful

---
 llvm/lib/Target/X86/X86.h                     |   2 +
 llvm/lib/Target/X86/X86MatchJumptablePass.cpp | 504 +++++++++---------
 llvm/lib/Target/X86/X86MatchJumptablePass.h   |  46 +-
 3 files changed, 308 insertions(+), 244 deletions(-)

diff --git a/llvm/lib/Target/X86/X86.h b/llvm/lib/Target/X86/X86.h
index 48a3fe1934a967..5d4882b921a23c 100644
--- a/llvm/lib/Target/X86/X86.h
+++ b/llvm/lib/Target/X86/X86.h
@@ -54,6 +54,8 @@ FunctionPass *createX86IndirectBranchTrackingPass();
 /// This will prevent a stall when returning on the Atom.
 FunctionPass *createX86PadShortFunctions();
 
+FunctionPass *createX86MatchJumptablePass();
+
 /// Return a pass that selectively replaces certain instructions (like add,
 /// sub, inc, dec, some shifts, and some multiplies) by equivalent LEA
 /// instructions, in order to eliminate execution delays in some processors.
diff --git a/llvm/lib/Target/X86/X86MatchJumptablePass.cpp b/llvm/lib/Target/X86/X86MatchJumptablePass.cpp
index c43d7c5b676b49..cf2a8ffe56345f 100644
--- a/llvm/lib/Target/X86/X86MatchJumptablePass.cpp
+++ b/llvm/lib/Target/X86/X86MatchJumptablePass.cpp
@@ -1,290 +1,316 @@
-// #include "X86.h"
-// #include "llvm/CodeGen/MachineFunctionPass.h"
-// #include "llvm/CodeGen/MachineInstr.h"
-// #include "llvm/CodeGen/MachineRegisterInfo.h"
-// #include "llvm/Support/Debug.h"
-// #include "llvm/Support/raw_ostream.h"
-// #include "X86MatchJumptablePass.h"
-// #include "llvm/CodeGen/MachineJumpTableInfo.h"
-
-// #define DEBUG_TYPE "match-jump-table"
-
-// using namespace llvm;
-
-// namespace {
-//   class X86MatchJumptablePass : public MachineFunctionPass {
-//   public:
-//     static char ID;
-
-//     X86MatchJumptablePass() : MachineFunctionPass(ID) {}
-
-//     bool runOnMachineFunction(MachineFunction &MF) override {
-//     LLVM_DEBUG(dbgs() << "Analyzing jump tables in function: " << MF.getName() << "\n");
-
-//     // Get jump table information
-//     MachineJumpTableInfo *JumpTableInfo = MF.getJumpTableInfo();
-//     if (!JumpTableInfo) {
-//       LLVM_DEBUG(dbgs() << "No jump tables in this function.\n");
-//       return false;
-//     }
-//     // Assuming JumpTableInfo is available
-//     for (unsigned JTIndex = 0; JTIndex < JumpTableInfo->getJumpTables().size(); ++JTIndex) {
-//   const MachineJumpTableEntry &JTEntry = JumpTableInfo->getJumpTables()[JTIndex];
-  
-//   LLVM_DEBUG(dbgs() << "Jump Table #" << JTIndex << " Base Address: " << JTEntry.BaseAddress << "\n");
-
-//   // Iterate through the entries (target basic blocks) in this jump table
-//   for (auto *MBB : JTEntry.MBBs) {
-//     LLVM_DEBUG(dbgs() << "  Target BasicBlock: " << MBB->getName() << " Address: " << MBB->getAddress() << "\n");
-//      }
-    
-
-//       // Trace potential indirect jumps related to this jump table
-//       traceIndirectJumps(MF, JTIndex, JumpTableInfo);
-//     }
-//     return false;
-
-// }
-
-//   void traceIndirectJumps(MachineFunction &MF, unsigned JTIndex, MachineJumpTableInfo *JumpTableInfo) {
-//     const MachineJumpTableEntry &JTEntry = JumpTableInfo->getJumpTables()[JTIndex];
-
-//     for (auto &MBB : MF) {
-//       for (auto &MI : MBB) {
-//         if (MI.isIndirectBranch()) {
-//           LLVM_DEBUG(dbgs() << "Found indirect jump: " << MI << "\n");
-
-//           // Analyze data flow to check if this jump is related to the jump table
-//           if (isJumpTableRelated(MI, JTEntry, MF)) {
-//             LLVM_DEBUG(dbgs() << "This indirect jump is related to Jump Table #" << JTIndex << "\n");
-//           }
-//         }
-//       }
-//     }
-//   }
-
-//   bool isJumpTableRelated(MachineInstr &MI, const MachineJumpTableEntry &JTEntry, MachineFunction &MF) {
-//   for (unsigned OpIdx = 0; OpIdx < MI.getNumOperands(); ++OpIdx) {
-//     const MachineOperand &Op = MI.getOperand(OpIdx);
-
-//     if (Op.isReg()) {
-//       Register Reg = Op.getReg();
-//       MachineRegisterInfo &MRI = MF.getRegInfo();
-//       // Check if any of the definitions of the register are related to a jump table load
-//       for (MachineInstr &DefMI : MRI.def_instructions(Reg)) {
-//         if (isJumpTableLoad(DefMI, JTEntry)) {
-//           return true;
-//         }
-//       }
-//     } else if (Op.isImm()) {
-//       // Check if the immediate operand might be an offset/index into the jump table
-//       int64_t ImmValue = Op.getImm();
-      
-//       // For example, if the jump table has 10 entries, check if the immediate is between 0 and 9
-//       if (ImmValue >= 0 && ImmValue < JTEntry.MBBs.size()) {
-//         // This immediate value could be an index into the jump table
-//         LLVM_DEBUG(dbgs() << "Immediate operand is a possible jump table index: " << ImmValue << "\n");
-//         return true;
-//       }
-//     }
-//   }
-//   return false;
-// }
-
-// bool isJumpTableLoad(MachineInstr &MI, const MachineJumpTableEntry &JTEntry) {
-//     if (MI.mayLoad()) {
-//         for (unsigned i = 0; i < MI.getNumOperands(); ++i) {
-//             const MachineOperand &Op = MI.getOperand(i);
-//             if (Op.isGlobal() && Op.getGlobal() == JTEntry.BaseAddress) {
-//                 return true;
-//             }
-//         }
-//     }
-//     return false;
-// }
-
-//   StringRef getPassName() const override {
-//     return "Match Jump Table Pass";
-//   }
-
-//     // StringRef getPassName() const override {
-//     //   return "X86 My Backend Pass";
-//     // }
-//   };
-// }
-
-// char X86MatchJumptablePass::ID = 0;
-
-// // Ensure the function is in the llvm namespace
-// namespace llvm {
-  
-//   // Define the pass
-//   FunctionPass *createX86MatchJumptablePass() {
-//     return new X86MatchJumptablePass();
-//   }
-
-// } // end llvm namespace
-
-// static RegisterPass<X86MatchJumptablePass> X("match-jump-table", "Match Jump Table Pass", false, false);
-
-
+// X86MatchJumptablePass.cpp
+#include "X86MatchJumptablePass.h"
 #include "X86.h"
+#include "X86InstrInfo.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
-#include "llvm/CodeGen/MachineInstr.h"
-#include "llvm/CodeGen/MachineRegisterInfo.h"
-#include "llvm/Support/Debug.h"
-#include "llvm/Support/raw_ostream.h"
-#include "llvm/CodeGen/MachineJumpTableInfo.h"
+#include "llvm/CodeGen/TargetInstrInfo.h"
+#include "llvm/PassRegistry.h"
+#include "llvm/Pass.h"
 
 #define DEBUG_TYPE "match-jump-table"
 
-using namespace llvm;
+namespace llvm {
 
-namespace {
-class X86MatchJumptablePass : public MachineFunctionPass {
-public:
-  static char ID;
+char X86MatchJumptablePass::ID = 0;
 
-  X86MatchJumptablePass() : MachineFunctionPass(ID) {}
+void initializeX86MatchJumptablePassPass(PassRegistry &Registry) {
+  RegisterPass<X86MatchJumptablePass> X("match-jump-table", "Match Jump Table Pass", false, false);
+}
 
-  bool runOnMachineFunction(MachineFunction &MF) override {
+void X86MatchJumptablePass::insertIdentifyingMarker(MachineInstr* MI, 
+                                                   MachineFunction &MF, 
+                                                   unsigned JTIndex) {
+  // Create a unique identifier for this indirect jump
+  std::string marker = "IJUMP_" + MF.getName().str() + "_" + std::to_string(JTIndex);
+  
+  // Get target instruction info
+  const TargetInstrInfo *TII = MF.getSubtarget().getInstrInfo();
+  
+  // Insert a special no-op instruction before the indirect jump with debug info
+  BuildMI(*MI->getParent(), MI, MI->getDebugLoc(), 
+          TII->get(X86::NOOP))
+          .addExternalSymbol(marker.c_str());
+}
+
+void X86MatchJumptablePass::recordJumpLocation(MachineInstr* MI, 
+                                              MachineFunction &MF, 
+                                              unsigned JTIndex) {
+  // Create a new section entry
+  auto &Context = MF.getFunction().getContext();
+  auto *M = MF.getFunction().getParent();
+  
+  // Create mapping data
+  std::string sectionName = ".jumptable.map";
+  std::string data = MF.getName().str() + "," + 
+                    std::to_string(JTIndex) + "\n";
+  
+  // Add the data to a special section
+  auto *GV = new GlobalVariable(
+    *M,
+    ArrayType::get(Type::getInt8Ty(Context), data.size() + 1),
+    true,
+    GlobalValue::PrivateLinkage,
+    ConstantDataArray::getString(Context, data),
+    "jump_map_entry"
+  );
+  
+  GV->setSection(sectionName);
+}
+
+bool X86MatchJumptablePass::runOnMachineFunction(MachineFunction &MF) {
+    static int PassRunCount = 0;
+    PassRunCount++;
+    
+    LLVM_DEBUG(dbgs() << "\n=== Pass Run #" << PassRunCount << " ===\n");
+    LLVM_DEBUG(dbgs() << "Current optimization phase: " 
+               << MF.getFunction().getParent()->getSourceFileName() << "\n");
     LLVM_DEBUG(dbgs() << "Analyzing jump tables in function: " << MF.getName() << "\n");
 
     // Get jump table information
     MachineJumpTableInfo *JumpTableInfo = MF.getJumpTableInfo();
     if (!JumpTableInfo) {
-      LLVM_DEBUG(dbgs() << "No jump tables in this function.\n");
-      return false;
+        LLVM_DEBUG(dbgs() << "No jump tables in this function.\n");
+        return false;
     }
 
+    bool Modified = false;
     for (unsigned JTIndex = 0; JTIndex < JumpTableInfo->getJumpTables().size(); ++JTIndex) {
-      const MachineJumpTableEntry &JTEntry = JumpTableInfo->getJumpTables()[JTIndex];
+        const MachineJumpTableEntry &JTEntry = JumpTableInfo->getJumpTables()[JTIndex];
+
+        LLVM_DEBUG(dbgs() << "Jump Table #" << JTIndex << " contains " 
+                         << JTEntry.MBBs.size() << " entries.\n");
+
+        // Print detailed information about jump table entries
+        LLVM_DEBUG(dbgs() << "Jump table entries:\n");
+        for (unsigned i = 0; i < JTEntry.MBBs.size(); ++i) {
+            if (JTEntry.MBBs[i]) {
+                LLVM_DEBUG(dbgs() << "  [" << i << "] -> BB#" 
+                           << JTEntry.MBBs[i]->getNumber() 
+                           << " (" << JTEntry.MBBs[i]->getName() << ")\n");
+            }
+        }
 
-      LLVM_DEBUG(dbgs() << "Jump Table #" << JTIndex << " contains " << JTEntry.MBBs.size()
-                        << " entries.\n");
+        // Find and process indirect jumps
+        MachineInstr* indirectJumpInstr = traceIndirectJumps(MF, JTIndex, JumpTableInfo);
+
+        if (indirectJumpInstr) {
+            // Print more detailed information about the indirect jump
+            LLVM_DEBUG(dbgs() << "Found indirect jump: " << *indirectJumpInstr << "\n");
+            LLVM_DEBUG(dbgs() << "Parent basic block: " << indirectJumpInstr->getParent()->getName() << "\n");
+            
+            // Print operand information
+            LLVM_DEBUG(dbgs() << "Jump instruction operands:\n");
+            for (unsigned i = 0; i < indirectJumpInstr->getNumOperands(); ++i) {
+                LLVM_DEBUG(dbgs() << "  Operand " << i << ": " << indirectJumpInstr->getOperand(i) << "\n");
+            }
 
-      // Iterate through the entries (target basic blocks) in this jump table
-      for (auto *MBB : JTEntry.MBBs) {
-        if (MBB) {
-          LLVM_DEBUG(dbgs() << "  Target BasicBlock: " << MBB->getName() << "\n");
-        }
-      }
-
-      // Assuming you have access to MF, JTIndex, and JumpTableInfo
-    MachineInstr* indirectJumpInstr = traceIndirectJumps(MF, JTIndex, JumpTableInfo);
-
-      if (indirectJumpInstr) {
-          // Handle the found indirect jump instruction
-          dbgs() << "Indirect jump found at address: " << indirectJumpInstr << "\n";
-          for (auto &MBB : JTEntry.MBBs) {
-          LLVM_DEBUG(dbgs() << "Address of MBB: " << &MBB << "\n"); // Print the address of the MBB
-          // Optionally print the name and instructions inside the MBB
-         // LLVM_DEBUG(dbgs() << "MBB: " << MBB.getName() << "\n");
-          // for (auto &MI : MBB) {
-          //     LLVM_DEBUG(dbgs() << "  Instruction: " << MI << "\n");
-          // }
+            // Insert marker and record location
+            insertIdentifyingMarker(indirectJumpInstr, MF, JTIndex);
+            recordJumpLocation(indirectJumpInstr, MF, JTIndex);
+            Modified = true;
         }
-      } 
     }
 
-    return false;
-  }
-
-  // void traceIndirectJumps(MachineFunction &MF, unsigned JTIndex, MachineJumpTableInfo *JumpTableInfo) {
-  //   const MachineJumpTableEntry &JTEntry = JumpTableInfo->getJumpTables()[JTIndex];
-
-  //   for (auto &MBB : MF) {
-  //     for (auto &MI : MBB) {
-  //       if (MI.isIndirectBranch()) {
-  //         LLVM_DEBUG(dbgs() << "Found indirect jump: " << MI << "\n");
-
-  //         // Analyze data flow to check if this jump is related to the jump table
-  //         if (isJumpTableRelated(MI, JTEntry, MF)) {
-  //           LLVM_DEBUG(dbgs() << "This indirect jump is related to Jump Table #" << JTIndex << "\n");
-  //         }
-  //       }
-  //     }
-  //   }
-  // }
-  MachineInstr* traceIndirectJumps(MachineFunction &MF, unsigned JTIndex, MachineJumpTableInfo *JumpTableInfo) {
+    return Modified;
+}
+
+MachineInstr* X86MatchJumptablePass::traceIndirectJumps(MachineFunction &MF, 
+                                                        unsigned JTIndex,
+                                                        MachineJumpTableInfo *JumpTableInfo) {
     const MachineJumpTableEntry &JTEntry = JumpTableInfo->getJumpTables()[JTIndex];
 
+    LLVM_DEBUG(dbgs() << "Tracing indirect jumps:\n");
     for (auto &MBB : MF) {
+        LLVM_DEBUG(dbgs() << "  Checking BB: " << MBB.getName() << "\n");
         for (auto &MI : MBB) {
+            LLVM_DEBUG(dbgs() << "    Checking instruction: " << MI << "\n");
             if (MI.isIndirectBranch()) {
-                LLVM_DEBUG(dbgs() << "Found indirect jump: " << MI << "\n");
+                LLVM_DEBUG(dbgs() << "    Found indirect jump: " << MI << "\n");
 
-                // Analyze data flow to check if this jump is related to the jump table
                 if (isJumpTableRelated(MI, JTEntry, MF)) {
-                    LLVM_DEBUG(dbgs() << "This indirect jump is related to Jump Table #" << JTIndex << "\n");
-                    
-                    // Return the address of the indirect jump (MI)
+                    LLVM_DEBUG(dbgs() << "    This indirect jump is related to Jump Table #" 
+                              << JTIndex << "\n");
                     return &MI;
+                } else {
+                    LLVM_DEBUG(dbgs() << "    Jump is not related to this jump table\n");
                 }
             }
         }
     }
 
-    // Return nullptr if no indirect jump is found
+    LLVM_DEBUG(dbgs() << "  No related indirect jump found\n");
     return nullptr;
 }
 
-  bool isJumpTableRelated(MachineInstr &MI, const MachineJumpTableEntry &JTEntry, MachineFunction &MF) {
-    for (unsigned OpIdx = 0; OpIdx < MI.getNumOperands(); ++OpIdx) {
-      const MachineOperand &Op = MI.getOperand(OpIdx);
+bool X86MatchJumptablePass::isJumpTableLoad(MachineInstr &MI, const MachineJumpTableEntry &JTEntry) {
+    LLVM_DEBUG(dbgs() << "\nAnalyzing potential jump table load instruction: " << MI << "\n");
+
+    // First check memory operands for jump table metadata
+    for (const MachineMemOperand *MMO : MI.memoperands()) {
+        LLVM_DEBUG(dbgs() << "  Checking memory operand flags: " << MMO->getFlags() << "\n");
+        if (MMO->getValue()) {
+            StringRef ValueName = MMO->getValue()->getName();
+            LLVM_DEBUG(dbgs() << "    Memory value name: '" << ValueName << "'\n");
+            if (ValueName.contains("jump-table")) {
+                LLVM_DEBUG(dbgs() << "    Found jump table in memory value name\n");
+                return true;
+            }
+        }
 
-      if (Op.isReg()) {
-        Register Reg = Op.getReg();
-        MachineRegisterInfo &MRI = MF.getRegInfo();
-        // Check if any of the definitions of the register are related to a jump table load
-        for (MachineInstr &DefMI : MRI.def_instructions(Reg)) {
-          if (isJumpTableLoad(DefMI, JTEntry)) {
-            return true;
-          }
+        // Check if this is a jump table load directly from memory operand comments
+        if (MI.getDesc().mayLoad() && MI.hasOneMemOperand()) {
+            // Look for jump table reference in the instruction's debug info or comments
+            if (MI.getDebugLoc()) {
+                std::string Comment;
+                raw_string_ostream OS(Comment);
+                MI.print(OS);
+                if (Comment.find("jump-table") != std::string::npos) {
+                    LLVM_DEBUG(dbgs() << "    Found jump table reference in instruction comment\n");
+                    return true;
+                }
+            }
+        }
+    }
+
+    // Check for the MOVSX pattern
+    if (MI.getOpcode() == X86::MOVSX64rm32) {
+        LLVM_DEBUG(dbgs() << "  Found MOVSX64rm32 instruction\n");
+        Register BaseReg;
+        
+        // Find base register
+        for (const MachineOperand &MO : MI.operands()) {
+            if (MO.isReg() && MO.isUse()) {
+                BaseReg = MO.getReg();
+                LLVM_DEBUG(dbgs() << "    Found base register: " << printReg(BaseReg, nullptr) << "\n");
+                break;
+            }
         }
-      } else if (Op.isImm()) {
-        // Check if the immediate operand might be an offset/index into the jump table
-        int64_t ImmValue = Op.getImm();
-        if (ImmValue >= 0 && ImmValue < JTEntry.MBBs.size()) {
-          // This immediate value could be an index into the jump table
-          LLVM_DEBUG(dbgs() << "Immediate operand is a possible jump table index: " << ImmValue
-                            << "\n");
-          return true;
+
+        if (BaseReg) {
+            // Look for preceding LEA
+            MachineBasicBlock::iterator MBBI = MI;
+            const MachineBasicBlock *MBB = MI.getParent();
+            
+            LLVM_DEBUG(dbgs() << "    Looking for LEA defining register: " << printReg(BaseReg, nullptr) << "\n");
+            
+            while (MBBI != MBB->begin()) {
+                --MBBI;
+                LLVM_DEBUG(dbgs() << "      Checking: " << *MBBI << "\n");
+                
+                if (MBBI->getOpcode() == X86::LEA64r) {
+                    LLVM_DEBUG(dbgs() << "      Found LEA64r\n");
+                    
+                    // Verify this LEA defines our base register
+                    const MachineOperand &DefReg = MBBI->getOperand(0);
+                    if (!DefReg.isReg() || DefReg.getReg() != BaseReg) {
+                        LLVM_DEBUG(dbgs() << "      LEA defines different register\n");
+                        continue;
+                    }
+                    
+                    // Check for jump table symbol
+                    for (const MachineOperand &MO : MBBI->operands()) {
+                        if (MO.isSymbol()) {
+                            StringRef SymName = MO.getSymbolName();
+                            LLVM_DEBUG(dbgs() << "      Checking symbol: '" << SymName << "'\n");
+                            if (SymName.contains("jump-table")) {
+                                LLVM_DEBUG(dbgs() << "      Found jump table symbol!\n");
+                                return true;
+                            }
+                        }
+                    }
+                }
+            }
+            LLVM_DEBUG(dbgs() << "    No matching LEA found\n");
         }
-      }
     }
+
     return false;
-  }
-
-  bool isJumpTableLoad(MachineInstr &MI, const MachineJumpTableEntry &JTEntry) {
-    if (MI.mayLoad()) {
-      for (unsigned i = 0; i < MI.getNumOperands(); ++i) {
-        const MachineOperand &Op = MI.getOperand(i);
-          if (Op.getType() == MachineOperand::MO_JumpTableIndex) {
-          unsigned JTIndex = Op.getIndex();
-          if (JTIndex < JTEntry.MBBs.size()) {
-            LLVM_DEBUG(dbgs() << "Instruction loads from Jump Table index: " << JTIndex << "\n");
-            return true;
-          }
+}
+
+bool X86MatchJumptablePass::isJumpTableRelated(MachineInstr &MI, 
+                                              const MachineJumpTableEntry &JTEntry,
+                                              MachineFunction &MF) {
+    if (!MI.isIndirectBranch()) {
+        LLVM_DEBUG(dbgs() << "Not an indirect branch, skipping\n");
+        return false;
+    }
+
+    LLVM_DEBUG(dbgs() << "\nAnalyzing indirect jump: " << MI << "\n");
+
+    // Get jump register
+    Register JumpReg;
+    for (const MachineOperand &MO : MI.operands()) {
+        if (MO.isReg() && MO.isUse()) {
+            JumpReg = MO.getReg();
+            LLVM_DEBUG(dbgs() << "Found jump register: " << printReg(JumpReg, nullptr) << "\n");
+            break;
         }
-      }
     }
-    return false;
-  }
 
-  StringRef getPassName() const override { return "Match Jump Table Pass"; }
-};
+    if (!JumpReg) {
+        LLVM_DEBUG(dbgs() << "No jump register found\n");
+        return false;
+    }
 
-} // namespace
+    SmallVector<MachineInstr*, 8> Worklist;
+    SmallPtrSet<MachineInstr*, 16> Visited;
+    
+    LLVM_DEBUG(dbgs() << "Starting backward analysis from register " << printReg(JumpReg, nullptr) << "\n");
 
-char X86MatchJumptablePass::ID = 0;
+    for (MachineInstr &DefMI : MF.getRegInfo().def_instructions(JumpReg)) {
+        Worklist.push_back(&DefMI);
+        LLVM_DEBUG(dbgs() << "Added to worklist: " << DefMI << "\n");
+    }
 
-namespace llvm {
+    while (!Worklist.empty()) {
+        MachineInstr *CurrMI = Worklist.pop_back_val();
+        if (!Visited.insert(CurrMI).second) {
+            LLVM_DEBUG(dbgs() << "Already visited: " << *CurrMI << "\n");
+            continue;
+        }
+
+        LLVM_DEBUG(dbgs() << "Analyzing instruction: " << *CurrMI << "\n");
+
+        if (isJumpTableLoad(*CurrMI, JTEntry)) {
+            LLVM_DEBUG(dbgs() << "Found jump table load!\n");
+            return true;
+        }
+
+        if (CurrMI->getOpcode() == X86::ADD64rr) {
+            LLVM_DEBUG(dbgs() << "Found ADD64rr, checking operands\n");
+            for (const MachineOperand &MO : CurrMI->operands()) {
+                if (MO.isReg() && MO.isUse()) {
+                    LLVM_DEBUG(dbgs() << "Checking register operand: " << printReg(MO.getReg(), nullptr) << "\n");
+                    for (MachineInstr &DefMI : MF.getRegInfo().def_instructions(MO.getReg())) {
+                        if (isJumpTableLoad(DefMI, JTEntry)) {
+                            LLVM_DEBUG(dbgs() << "Found jump table load via ADD operand!\n");
+                            return true;
+                        }
+                    }
+                }
+            }
+        }
 
-// Define the pass
-FunctionPass *createX86MatchJumptablePass() { return new X86MatchJumptablePass(); }
+        // Add uses to worklist
+        for (const MachineOperand &MO : CurrMI->operands()) {
+            if (MO.isReg() && MO.isUse()) {
+                LLVM_DEBUG(dbgs() << "Adding definitions of register " << printReg(MO.getReg(), nullptr) << " to worklist\n");
+                for (MachineInstr &DefMI : MF.getRegInfo().def_instructions(MO.getReg())) {
+                    if (!Visited.count(&DefMI)) {
+                        Worklist.push_back(&DefMI);
+                        LLVM_DEBUG(dbgs() << "Added to worklist: " << DefMI << "\n");
+                    }
+                }
+            }
+        }
+    }
 
-} // namespace llvm
+    LLVM_DEBUG(dbgs() << "No jump table relation found\n");
+    return false;
+}
+
+FunctionPass *createX86MatchJumptablePass() { 
+    return new X86MatchJumptablePass();
+}
 
-static RegisterPass<X86MatchJumptablePass> X("match-jump-table", "Match Jump Table Pass", false,
-                                             false);
+} // end namespace llvm
\ No newline at end of file
diff --git a/llvm/lib/Target/X86/X86MatchJumptablePass.h b/llvm/lib/Target/X86/X86MatchJumptablePass.h
index ac6fc315b025ad..779ada3d8b6d41 100644
--- a/llvm/lib/Target/X86/X86MatchJumptablePass.h
+++ b/llvm/lib/Target/X86/X86MatchJumptablePass.h
@@ -1,14 +1,50 @@
-#ifndef LLVM_LIB_TARGET_X86_X86MATCHJUMPTABLEPASS_H
-#define LLVM_LIB_TARGET_X86_X86MATCHJUMPTABLEPASS_H
+// X86MatchJumptablePass.h
+#ifndef X86_MATCH_JUMPTABLE_PASS_H
+#define X86_MATCH_JUMPTABLE_PASS_H
 
+#include "MCTargetDesc/X86MCTargetDesc.h"
+#include "X86.h"
+#include "X86InstrInfo.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/MachineInstr.h"
+#include "llvm/CodeGen/MachineRegisterInfo.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/raw_ostream.h"
+#include "llvm/CodeGen/MachineJumpTableInfo.h"
+#include "llvm/IR/GlobalVariable.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/CodeGen/TargetInstrInfo.h"
 
 namespace llvm {
 
-/// \brief Creates the X86MatchJumptablePass.
-/// This pass analyzes and processes jump tables in X86 backend code generation.
+class X86MatchJumptablePass : public MachineFunctionPass {
+private:
+  void insertIdentifyingMarker(MachineInstr* MI, MachineFunction &MF, unsigned JTIndex);
+  void recordJumpLocation(MachineInstr* MI, MachineFunction &MF, unsigned JTIndex);
+  MachineInstr* traceIndirectJumps(MachineFunction &MF, unsigned JTIndex, 
+                                  MachineJumpTableInfo *JumpTableInfo);
+  bool isJumpTableRelated(MachineInstr &MI, const MachineJumpTableEntry &JTEntry, 
+                         MachineFunction &MF);
+  bool isJumpTableLoad(MachineInstr &MI, const MachineJumpTableEntry &JTEntry);
+  bool isRegUsedInJumpTableLoad(Register Reg,MachineFunction &MF,
+                                                    const MachineJumpTableEntry &JTEntry);
+
+public:
+  static char ID;
+
+  X86MatchJumptablePass() : MachineFunctionPass(ID) {}
+
+  bool runOnMachineFunction(MachineFunction &MF) override;
+  StringRef getPassName() const override { return "Match Jump Table Pass"; }
+};
+
 FunctionPass *createX86MatchJumptablePass();
 
+// Pass initialization declaration
+void initializeX86MatchJumptablePassPass(PassRegistry &Registry);
+
 } // namespace llvm
 
-#endif // LLVM_LIB_TARGET_X86_X86MATCHJUMPTABLEPASS_H
+#endif // X86_MATCH_JUMPTABLE_PASS_H
+

>From 5e4c5bae0ef246aa0bc88d0d04ab13ea78fd7b36 Mon Sep 17 00:00:00 2001
From: ryan <liu864579887 at gmail.com>
Date: Mon, 13 Jan 2025 14:48:03 -0600
Subject: [PATCH 6/7] jump addr block insert and get successfully

---
 llvm/lib/Target/X86/X86MatchJumptablePass.cpp | 158 +++++++++++++++---
 1 file changed, 131 insertions(+), 27 deletions(-)

diff --git a/llvm/lib/Target/X86/X86MatchJumptablePass.cpp b/llvm/lib/Target/X86/X86MatchJumptablePass.cpp
index cf2a8ffe56345f..fddce219ee6d0e 100644
--- a/llvm/lib/Target/X86/X86MatchJumptablePass.cpp
+++ b/llvm/lib/Target/X86/X86MatchJumptablePass.cpp
@@ -57,16 +57,103 @@ void X86MatchJumptablePass::recordJumpLocation(MachineInstr* MI,
   GV->setSection(sectionName);
 }
 
+// bool X86MatchJumptablePass::runOnMachineFunction(MachineFunction &MF) {
+//     static int PassRunCount = 0;
+//     PassRunCount++;
+    
+//     LLVM_DEBUG(dbgs() << "\n=== Pass Run #" << PassRunCount << " ===\n");
+//     LLVM_DEBUG(dbgs() << "Current optimization phase: " 
+//                << MF.getFunction().getParent()->getSourceFileName() << "\n");
+//     LLVM_DEBUG(dbgs() << "Analyzing jump tables in function: " << MF.getName() << "\n");
+
+//     // Get jump table information
+//     MachineJumpTableInfo *JumpTableInfo = MF.getJumpTableInfo();
+//     if (!JumpTableInfo) {
+//         LLVM_DEBUG(dbgs() << "No jump tables in this function.\n");
+//         return false;
+//     }
+
+//     bool Modified = false;
+//     for (unsigned JTIndex = 0; JTIndex < JumpTableInfo->getJumpTables().size(); ++JTIndex) {
+//         const MachineJumpTableEntry &JTEntry = JumpTableInfo->getJumpTables()[JTIndex];
+
+//         LLVM_DEBUG(dbgs() << "Jump Table #" << JTIndex << " contains " 
+//                          << JTEntry.MBBs.size() << " entries.\n");
+
+//         // Print detailed information about jump table entries
+//         LLVM_DEBUG(dbgs() << "Jump table entries:\n");
+//         for (unsigned i = 0; i < JTEntry.MBBs.size(); ++i) {
+//             if (JTEntry.MBBs[i]) {
+//                 LLVM_DEBUG(dbgs() << "  [" << i << "] -> BB#" 
+//                            << JTEntry.MBBs[i]->getNumber() 
+//                            << " (" << JTEntry.MBBs[i]->getName() << ")\n");
+//             }
+//         }
+
+//         // Find and process indirect jumps
+//         MachineInstr* indirectJumpInstr = traceIndirectJumps(MF, JTIndex, JumpTableInfo);
+
+//         if (indirectJumpInstr) {
+//             LLVM_DEBUG(dbgs() << "Found indirect jump: " << *indirectJumpInstr << "\n");
+//                 std::string VarName = "ijump_addr_"  + std::to_string(JTIndex);
+//                 // Get the parent basic block
+//                 MachineBasicBlock *MBB = indirectJumpInstr->getParent();
+//                 // Get the offset of instruction within the block
+//             uint64_t Offset = 0;
+//             for (auto &MI : *MBB) {
+//                 if (&MI == indirectJumpInstr)
+//                     break;
+//                 Offset += MI.getDesc().getSize();
+//             }
+            
+//             // Now add this offset to the block address
+//             Constant *BlockAddr = BlockAddress::get(&MF.getFunction(), const_cast<BasicBlock*>(MBB->getBasicBlock()));
+//             Constant *OffsetConst = ConstantInt::get(Type::getInt64Ty(M->getContext()), Offset);
+//             Constant *AddrConst = ConstantExpr::getAdd(
+//                 ConstantExpr::getPtrToInt(BlockAddr, Type::getInt64Ty(M->getContext())),
+//                 OffsetConst
+//             );
+            
+                
+//                 auto *GV = new GlobalVariable(
+//                     *M,
+//                     Type::getInt64Ty(M->getContext()),
+//                     true,  // isConstant
+//                     GlobalValue::ExternalLinkage,
+//                     AddrConst,
+//                     VarName,
+//                     nullptr,                    // InsertBefore
+//                     GlobalValue::NotThreadLocal, // Thread Local
+//                     0,                          // AddressSpace
+//                     true                        // Constant
+//                 );
+//                 GV->setVisibility(GlobalValue::DefaultVisibility);
+//                 // Set the section
+//                 GV->setSection(".ijump.addr");
+//                 Modified = true;
+//         }
+//     }
+
+//     return Modified;
+// }
+
 bool X86MatchJumptablePass::runOnMachineFunction(MachineFunction &MF) {
     static int PassRunCount = 0;
     PassRunCount++;
     
     LLVM_DEBUG(dbgs() << "\n=== Pass Run #" << PassRunCount << " ===\n");
-    LLVM_DEBUG(dbgs() << "Current optimization phase: " 
-               << MF.getFunction().getParent()->getSourceFileName() << "\n");
-    LLVM_DEBUG(dbgs() << "Analyzing jump tables in function: " << MF.getName() << "\n");
+    LLVM_DEBUG(dbgs() << "Analyzing function: " << MF.getName() << "\n");
+
+    // Get the function's address as a base
+    Function &F = MF.getFunction();
+    Constant *FuncAddr = ConstantExpr::getPtrToInt(
+    ConstantExpr::getBitCast(&F, PointerType::get(Type::getInt8Ty(F.getContext()), 0)),
+    Type::getInt64Ty(F.getContext())
+);
 
-    // Get jump table information
+    LLVM_DEBUG(dbgs() << "Function address: " << *FuncAddr << "\n");
+
+    // Process jump tables
     MachineJumpTableInfo *JumpTableInfo = MF.getJumpTableInfo();
     if (!JumpTableInfo) {
         LLVM_DEBUG(dbgs() << "No jump tables in this function.\n");
@@ -78,35 +165,51 @@ bool X86MatchJumptablePass::runOnMachineFunction(MachineFunction &MF) {
         const MachineJumpTableEntry &JTEntry = JumpTableInfo->getJumpTables()[JTIndex];
 
         LLVM_DEBUG(dbgs() << "Jump Table #" << JTIndex << " contains " 
-                         << JTEntry.MBBs.size() << " entries.\n");
-
-        // Print detailed information about jump table entries
-        LLVM_DEBUG(dbgs() << "Jump table entries:\n");
-        for (unsigned i = 0; i < JTEntry.MBBs.size(); ++i) {
-            if (JTEntry.MBBs[i]) {
-                LLVM_DEBUG(dbgs() << "  [" << i << "] -> BB#" 
-                           << JTEntry.MBBs[i]->getNumber() 
-                           << " (" << JTEntry.MBBs[i]->getName() << ")\n");
-            }
-        }
+                          << JTEntry.MBBs.size() << " entries.\n");
 
-        // Find and process indirect jumps
-        MachineInstr* indirectJumpInstr = traceIndirectJumps(MF, JTIndex, JumpTableInfo);
+        MachineInstr *indirectJumpInstr = traceIndirectJumps(MF, JTIndex, JumpTableInfo);
 
         if (indirectJumpInstr) {
-            // Print more detailed information about the indirect jump
             LLVM_DEBUG(dbgs() << "Found indirect jump: " << *indirectJumpInstr << "\n");
-            LLVM_DEBUG(dbgs() << "Parent basic block: " << indirectJumpInstr->getParent()->getName() << "\n");
-            
-            // Print operand information
-            LLVM_DEBUG(dbgs() << "Jump instruction operands:\n");
-            for (unsigned i = 0; i < indirectJumpInstr->getNumOperands(); ++i) {
-                LLVM_DEBUG(dbgs() << "  Operand " << i << ": " << indirectJumpInstr->getOperand(i) << "\n");
+
+            std::string VarName = "ijump_addr_" + std::to_string(JTIndex);
+
+            // Calculate the offset of the indirect jump from the function's start
+            uint64_t Offset = 0;
+            for (auto &MBB : MF) {
+                for (auto &MI : MBB) {
+                    if (&MI == indirectJumpInstr)
+                        break;
+                    Offset += MI.getDesc().getSize();
+                }
+                if (Offset > 0) break; // Stop after finding the indirect jump
             }
 
-            // Insert marker and record location
-            insertIdentifyingMarker(indirectJumpInstr, MF, JTIndex);
-            recordJumpLocation(indirectJumpInstr, MF, JTIndex);
+            LLVM_DEBUG(dbgs() << "Offset from function base: " << Offset << "\n");
+
+            // Add the offset to the function address
+            Constant *OffsetConst = ConstantInt::get(Type::getInt64Ty(F.getContext()), Offset);
+            Constant *AddrConst = ConstantExpr::getAdd(FuncAddr, OffsetConst);
+
+            // Emit the address as a global variable
+            auto *GV = new GlobalVariable(
+                *F.getParent(),
+                Type::getInt64Ty(F.getContext()),
+                true,  // isConstant
+                GlobalValue::ExternalLinkage,
+                AddrConst,
+                VarName,
+                nullptr,                    // InsertBefore
+                GlobalValue::NotThreadLocal, // Thread Local
+                0,                          // AddressSpace
+                true                        // Constant
+            );
+            GV->setVisibility(GlobalValue::DefaultVisibility);
+            GV->setSection(".ijump.addr");
+
+            LLVM_DEBUG(dbgs() << "Emitted symbol: " << VarName 
+                              << " -> Address: " << *AddrConst << "\n");
+
             Modified = true;
         }
     }
@@ -114,6 +217,7 @@ bool X86MatchJumptablePass::runOnMachineFunction(MachineFunction &MF) {
     return Modified;
 }
 
+
 MachineInstr* X86MatchJumptablePass::traceIndirectJumps(MachineFunction &MF, 
                                                         unsigned JTIndex,
                                                         MachineJumpTableInfo *JumpTableInfo) {

>From 637d4b61c2596b85e3690a48db77ece478c00f06 Mon Sep 17 00:00:00 2001
From: ryan <liu864579887 at gmail.com>
Date: Thu, 16 Jan 2025 23:22:56 -0600
Subject: [PATCH 7/7] jumptable successful

---
 debug.log                                     |   2 +
 llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp    |  30 ++-
 llvm/lib/Target/X86/X86MatchJumptablePass.cpp | 228 ++++--------------
 llvm/lib/Target/X86/X86MatchJumptablePass.h   |   2 -
 4 files changed, 70 insertions(+), 192 deletions(-)
 create mode 100644 debug.log

diff --git a/debug.log b/debug.log
new file mode 100644
index 00000000000000..b0a65a1bd5c972
--- /dev/null
+++ b/debug.log
@@ -0,0 +1,2 @@
+clang++: error: no such file or directory: 'test_jump_table.cpp'
+clang++: error: no input files
diff --git a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
index 3ba45900e45691..a4f3a94d5aff67 100644
--- a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
+++ b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
@@ -2862,6 +2862,7 @@ void AsmPrinter::emitJumpTableInfo() {
   // Pick the directive to use to print the jump table entries, and switch to
   // the appropriate section.
   const Function &F = MF->getFunction();
+  // errs() << "F" << F.getName() << "\n";
   const TargetLoweringObjectFile &TLOF = getObjFileLowering();
   bool JTInDiffSection = !TLOF.shouldPutJumpTableInFunctionSection(
       MJTI->getEntryKind() == MachineJumpTableInfo::EK_LabelDifference32 ||
@@ -2882,27 +2883,34 @@ void AsmPrinter::emitJumpTableInfo() {
 
   for (unsigned JTI = 0, e = JT.size(); JTI != e; ++JTI) {
     const std::vector<MachineBasicBlock*> &JTBBs = JT[JTI].MBBs;
+    
 
     // If this jump table was deleted, ignore it.
     if (JTBBs.empty()) continue;
+    // errs() <<"test0" << "\n";
 
     // For the EK_LabelDifference32 entry, if using .set avoids a relocation,
     /// emit a .set directive for each unique entry.
     if (MJTI->getEntryKind() == MachineJumpTableInfo::EK_LabelDifference32 &&
         MAI->doesSetDirectiveSuppressReloc()) {
+      //  errs() << "test1" <<"\n";
       SmallPtrSet<const MachineBasicBlock*, 16> EmittedSets;
       const TargetLowering *TLI = MF->getSubtarget().getTargetLowering();
       const MCExpr *Base = TLI->getPICJumpTableRelocBaseExpr(MF,JTI,OutContext);
       for (const MachineBasicBlock *MBB : JTBBs) {
         if (!EmittedSets.insert(MBB).second)
           continue;
-
+        // errs() <<"test2" <<"\n";
         // .set LJTSet, LBB32-base
+        MCSymbol *MCsy = MBB->getSymbol();
+        MCSymbol *JTsetsy = GetJTSetSymbol(JTI, MBB->getNumber());
         const MCExpr *LHS =
-          MCSymbolRefExpr::create(MBB->getSymbol(), OutContext);
-        OutStreamer->emitAssignment(GetJTSetSymbol(JTI, MBB->getNumber()),
+          MCSymbolRefExpr::create(MCsy, OutContext);
+        OutStreamer->emitAssignment(JTsetsy,
                                     MCBinaryExpr::createSub(LHS, Base,
                                                             OutContext));
+          // errs() << "Symbol" << MCsy->getName() << "\n";
+          // errs() << "JTSymbol" << JTsetsy->getName() << "\n";
       }
     }
 
@@ -2921,8 +2929,10 @@ void AsmPrinter::emitJumpTableInfo() {
 
     // Defer MCAssembler based constant folding due to a performance issue. The
     // label differences will be evaluated at write time.
-    for (const MachineBasicBlock *MBB : JTBBs)
+    for (const MachineBasicBlock *MBB : JTBBs){
+      // errs() <<"test6" << "\n";
       emitJumpTableEntry(MJTI, MBB, JTI);
+    }
   }
 
   if (EmitJumpTableSizesSection)
@@ -2987,6 +2997,7 @@ void AsmPrinter::emitJumpTableEntry(const MachineJumpTableInfo *MJTI,
                                     unsigned UID) const {
   assert(MBB && MBB->getNumber() >= 0 && "Invalid basic block");
   const MCExpr *Value = nullptr;
+  // errs() << "EntryKind:" << MJTI->getEntryKind() <<"\n";
   switch (MJTI->getEntryKind()) {
   case MachineJumpTableInfo::EK_Inline:
     llvm_unreachable("Cannot emit EK_Inline jump table entry");
@@ -3026,16 +3037,25 @@ void AsmPrinter::emitJumpTableEntry(const MachineJumpTableInfo *MJTI,
     // If the .set directive avoids relocations, this is emitted as:
     //      .set L4_5_set_123, LBB123 - LJTI1_2
     //      .word L4_5_set_123
+      MCSymbol *JTsetsy = GetJTSetSymbol(UID, MBB->getNumber());
     if (MJTI->getEntryKind() == MachineJumpTableInfo::EK_LabelDifference32 &&
         MAI->doesSetDirectiveSuppressReloc()) {
-      Value = MCSymbolRefExpr::create(GetJTSetSymbol(UID, MBB->getNumber()),
+      Value = MCSymbolRefExpr::create(JTsetsy,
                                       OutContext);
+      // errs() <<"JTSetsy1:" <<JTsetsy->getName() <<"\n";
+      errs() << "kind1:" <<Value->getKind() <<"\n";
+      Value->dump();
+      
       break;
     }
     Value = MCSymbolRefExpr::create(MBB->getSymbol(), OutContext);
     const TargetLowering *TLI = MF->getSubtarget().getTargetLowering();
     const MCExpr *Base = TLI->getPICJumpTableRelocBaseExpr(MF, UID, OutContext);
     Value = MCBinaryExpr::createSub(Value, Base, OutContext);
+      errs() <<"Base kind" <<(int)Base->getKind() <<"\n";
+      errs() << "Value kind" <<(int)Value->getKind() <<"\n";
+      Value->dump();
+      Base->dump();
     break;
   }
   }
diff --git a/llvm/lib/Target/X86/X86MatchJumptablePass.cpp b/llvm/lib/Target/X86/X86MatchJumptablePass.cpp
index fddce219ee6d0e..e4f8a948c413b1 100644
--- a/llvm/lib/Target/X86/X86MatchJumptablePass.cpp
+++ b/llvm/lib/Target/X86/X86MatchJumptablePass.cpp
@@ -6,8 +6,11 @@
 #include "llvm/CodeGen/TargetInstrInfo.h"
 #include "llvm/PassRegistry.h"
 #include "llvm/Pass.h"
+#include "llvm/MC/MCContext.h"
+#include "llvm/MC/MCSymbol.h"
 
 #define DEBUG_TYPE "match-jump-table"
+ static int RunCount = 0;
 
 namespace llvm {
 
@@ -17,141 +20,15 @@ void initializeX86MatchJumptablePassPass(PassRegistry &Registry) {
   RegisterPass<X86MatchJumptablePass> X("match-jump-table", "Match Jump Table Pass", false, false);
 }
 
-void X86MatchJumptablePass::insertIdentifyingMarker(MachineInstr* MI, 
-                                                   MachineFunction &MF, 
-                                                   unsigned JTIndex) {
-  // Create a unique identifier for this indirect jump
-  std::string marker = "IJUMP_" + MF.getName().str() + "_" + std::to_string(JTIndex);
-  
-  // Get target instruction info
-  const TargetInstrInfo *TII = MF.getSubtarget().getInstrInfo();
-  
-  // Insert a special no-op instruction before the indirect jump with debug info
-  BuildMI(*MI->getParent(), MI, MI->getDebugLoc(), 
-          TII->get(X86::NOOP))
-          .addExternalSymbol(marker.c_str());
-}
-
-void X86MatchJumptablePass::recordJumpLocation(MachineInstr* MI, 
-                                              MachineFunction &MF, 
-                                              unsigned JTIndex) {
-  // Create a new section entry
-  auto &Context = MF.getFunction().getContext();
-  auto *M = MF.getFunction().getParent();
-  
-  // Create mapping data
-  std::string sectionName = ".jumptable.map";
-  std::string data = MF.getName().str() + "," + 
-                    std::to_string(JTIndex) + "\n";
-  
-  // Add the data to a special section
-  auto *GV = new GlobalVariable(
-    *M,
-    ArrayType::get(Type::getInt8Ty(Context), data.size() + 1),
-    true,
-    GlobalValue::PrivateLinkage,
-    ConstantDataArray::getString(Context, data),
-    "jump_map_entry"
-  );
-  
-  GV->setSection(sectionName);
-}
-
-// bool X86MatchJumptablePass::runOnMachineFunction(MachineFunction &MF) {
-//     static int PassRunCount = 0;
-//     PassRunCount++;
-    
-//     LLVM_DEBUG(dbgs() << "\n=== Pass Run #" << PassRunCount << " ===\n");
-//     LLVM_DEBUG(dbgs() << "Current optimization phase: " 
-//                << MF.getFunction().getParent()->getSourceFileName() << "\n");
-//     LLVM_DEBUG(dbgs() << "Analyzing jump tables in function: " << MF.getName() << "\n");
-
-//     // Get jump table information
-//     MachineJumpTableInfo *JumpTableInfo = MF.getJumpTableInfo();
-//     if (!JumpTableInfo) {
-//         LLVM_DEBUG(dbgs() << "No jump tables in this function.\n");
-//         return false;
-//     }
-
-//     bool Modified = false;
-//     for (unsigned JTIndex = 0; JTIndex < JumpTableInfo->getJumpTables().size(); ++JTIndex) {
-//         const MachineJumpTableEntry &JTEntry = JumpTableInfo->getJumpTables()[JTIndex];
-
-//         LLVM_DEBUG(dbgs() << "Jump Table #" << JTIndex << " contains " 
-//                          << JTEntry.MBBs.size() << " entries.\n");
-
-//         // Print detailed information about jump table entries
-//         LLVM_DEBUG(dbgs() << "Jump table entries:\n");
-//         for (unsigned i = 0; i < JTEntry.MBBs.size(); ++i) {
-//             if (JTEntry.MBBs[i]) {
-//                 LLVM_DEBUG(dbgs() << "  [" << i << "] -> BB#" 
-//                            << JTEntry.MBBs[i]->getNumber() 
-//                            << " (" << JTEntry.MBBs[i]->getName() << ")\n");
-//             }
-//         }
-
-//         // Find and process indirect jumps
-//         MachineInstr* indirectJumpInstr = traceIndirectJumps(MF, JTIndex, JumpTableInfo);
-
-//         if (indirectJumpInstr) {
-//             LLVM_DEBUG(dbgs() << "Found indirect jump: " << *indirectJumpInstr << "\n");
-//                 std::string VarName = "ijump_addr_"  + std::to_string(JTIndex);
-//                 // Get the parent basic block
-//                 MachineBasicBlock *MBB = indirectJumpInstr->getParent();
-//                 // Get the offset of instruction within the block
-//             uint64_t Offset = 0;
-//             for (auto &MI : *MBB) {
-//                 if (&MI == indirectJumpInstr)
-//                     break;
-//                 Offset += MI.getDesc().getSize();
-//             }
-            
-//             // Now add this offset to the block address
-//             Constant *BlockAddr = BlockAddress::get(&MF.getFunction(), const_cast<BasicBlock*>(MBB->getBasicBlock()));
-//             Constant *OffsetConst = ConstantInt::get(Type::getInt64Ty(M->getContext()), Offset);
-//             Constant *AddrConst = ConstantExpr::getAdd(
-//                 ConstantExpr::getPtrToInt(BlockAddr, Type::getInt64Ty(M->getContext())),
-//                 OffsetConst
-//             );
-            
-                
-//                 auto *GV = new GlobalVariable(
-//                     *M,
-//                     Type::getInt64Ty(M->getContext()),
-//                     true,  // isConstant
-//                     GlobalValue::ExternalLinkage,
-//                     AddrConst,
-//                     VarName,
-//                     nullptr,                    // InsertBefore
-//                     GlobalValue::NotThreadLocal, // Thread Local
-//                     0,                          // AddressSpace
-//                     true                        // Constant
-//                 );
-//                 GV->setVisibility(GlobalValue::DefaultVisibility);
-//                 // Set the section
-//                 GV->setSection(".ijump.addr");
-//                 Modified = true;
-//         }
-//     }
-
-//     return Modified;
-// }
-
 bool X86MatchJumptablePass::runOnMachineFunction(MachineFunction &MF) {
-    static int PassRunCount = 0;
-    PassRunCount++;
-    
-    LLVM_DEBUG(dbgs() << "\n=== Pass Run #" << PassRunCount << " ===\n");
-    LLVM_DEBUG(dbgs() << "Analyzing function: " << MF.getName() << "\n");
-
-    // Get the function's address as a base
+   
     Function &F = MF.getFunction();
     Constant *FuncAddr = ConstantExpr::getPtrToInt(
     ConstantExpr::getBitCast(&F, PointerType::get(Type::getInt8Ty(F.getContext()), 0)),
     Type::getInt64Ty(F.getContext())
 );
 
-    LLVM_DEBUG(dbgs() << "Function address: " << *FuncAddr << "\n");
+    LLVM_DEBUG(dbgs() << "Function address: " << FuncAddr << "\n");
 
     // Process jump tables
     MachineJumpTableInfo *JumpTableInfo = MF.getJumpTableInfo();
@@ -161,68 +38,49 @@ bool X86MatchJumptablePass::runOnMachineFunction(MachineFunction &MF) {
     }
 
     bool Modified = false;
+     LLVM_DEBUG(dbgs() << "Jump Table Size#" << JumpTableInfo->getJumpTables().size() << "\n");
+    
     for (unsigned JTIndex = 0; JTIndex < JumpTableInfo->getJumpTables().size(); ++JTIndex) {
-        const MachineJumpTableEntry &JTEntry = JumpTableInfo->getJumpTables()[JTIndex];
-
-        LLVM_DEBUG(dbgs() << "Jump Table #" << JTIndex << " contains " 
-                          << JTEntry.MBBs.size() << " entries.\n");
-
-        MachineInstr *indirectJumpInstr = traceIndirectJumps(MF, JTIndex, JumpTableInfo);
-
-        if (indirectJumpInstr) {
-            LLVM_DEBUG(dbgs() << "Found indirect jump: " << *indirectJumpInstr << "\n");
-
-            std::string VarName = "ijump_addr_" + std::to_string(JTIndex);
-
-            // Calculate the offset of the indirect jump from the function's start
-            uint64_t Offset = 0;
-            for (auto &MBB : MF) {
-                for (auto &MI : MBB) {
-                    if (&MI == indirectJumpInstr)
-                        break;
-                    Offset += MI.getDesc().getSize();
-                }
-                if (Offset > 0) break; // Stop after finding the indirect jump
-            }
-
-            LLVM_DEBUG(dbgs() << "Offset from function base: " << Offset << "\n");
-
-            // Add the offset to the function address
-            Constant *OffsetConst = ConstantInt::get(Type::getInt64Ty(F.getContext()), Offset);
-            Constant *AddrConst = ConstantExpr::getAdd(FuncAddr, OffsetConst);
-
-            // Emit the address as a global variable
-            auto *GV = new GlobalVariable(
-                *F.getParent(),
-                Type::getInt64Ty(F.getContext()),
-                true,  // isConstant
-                GlobalValue::ExternalLinkage,
-                AddrConst,
-                VarName,
-                nullptr,                    // InsertBefore
-                GlobalValue::NotThreadLocal, // Thread Local
-                0,                          // AddressSpace
-                true                        // Constant
-            );
-            GV->setVisibility(GlobalValue::DefaultVisibility);
-            GV->setSection(".ijump.addr");
-
-            LLVM_DEBUG(dbgs() << "Emitted symbol: " << VarName 
-                              << " -> Address: " << *AddrConst << "\n");
-
-            Modified = true;
+    const MachineJumpTableEntry &JTEntry = JumpTableInfo->getJumpTables()[JTIndex];
+    LLVM_DEBUG(dbgs() << "FuncAddr:" << FuncAddr << "Jump Table #" << JTIndex << " contains " 
+                      << JTEntry.MBBs.size() << " entries.\n");
+    
+    // Handle indirect jump instruction
+    MachineInstr *indirectJumpInstr = traceIndirectJumps(MF, JTIndex, JumpTableInfo);
+    if (indirectJumpInstr) {
+        // Create label for indirect jump
+        std::string LabelName =  std::to_string(RunCount) + "_IJUMP_" + std::to_string(JTIndex);
+        MCSymbol *Label = MF.getContext().getOrCreateSymbol(LabelName);
+        indirectJumpInstr->setPreInstrSymbol(MF, Label);
+        
+        for (unsigned EntryIndex = 0; EntryIndex < JTEntry.MBBs.size(); ++EntryIndex) {
+            MachineBasicBlock *TargetMBB = JTEntry.MBBs[EntryIndex];
+            if (!TargetMBB->empty()) {
+            std::string EntryLabelName = std::to_string(RunCount) + "_JTENTRY_" + std::to_string(JTIndex) + "_" + std::to_string(EntryIndex);
+            MCSymbol *EntryLabel = MF.getContext().getOrCreateSymbol(EntryLabelName);
+            
+            // Set label only on first instruction
+            MachineInstr &FirstInstr = TargetMBB->front();
+            FirstInstr.setPreInstrSymbol(MF, EntryLabel);
+            
+            LLVM_DEBUG(dbgs() << "Created label for jump table entry: " << EntryLabelName << "\n");
         }
+        }
+        RunCount ++;
+    }
+        Modified = true;
     }
 
+
     return Modified;
 }
 
 
-MachineInstr* X86MatchJumptablePass::traceIndirectJumps(MachineFunction &MF, 
-                                                        unsigned JTIndex,
-                                                        MachineJumpTableInfo *JumpTableInfo) {
-    const MachineJumpTableEntry &JTEntry = JumpTableInfo->getJumpTables()[JTIndex];
 
+MachineInstr* X86MatchJumptablePass::traceIndirectJumps(MachineFunction &MF,
+                                                          unsigned JTIndex,
+                                                          MachineJumpTableInfo *JumpTableInfo) {
+    const MachineJumpTableEntry &JTEntry = JumpTableInfo->getJumpTables()[JTIndex]; 
     LLVM_DEBUG(dbgs() << "Tracing indirect jumps:\n");
     for (auto &MBB : MF) {
         LLVM_DEBUG(dbgs() << "  Checking BB: " << MBB.getName() << "\n");
@@ -230,10 +88,10 @@ MachineInstr* X86MatchJumptablePass::traceIndirectJumps(MachineFunction &MF,
             LLVM_DEBUG(dbgs() << "    Checking instruction: " << MI << "\n");
             if (MI.isIndirectBranch()) {
                 LLVM_DEBUG(dbgs() << "    Found indirect jump: " << MI << "\n");
-
+                
                 if (isJumpTableRelated(MI, JTEntry, MF)) {
-                    LLVM_DEBUG(dbgs() << "    This indirect jump is related to Jump Table #" 
-                              << JTIndex << "\n");
+                    LLVM_DEBUG(dbgs() << "    This indirect jump is related to Jump Table #"
+                               << JTIndex << "\n");
                     return &MI;
                 } else {
                     LLVM_DEBUG(dbgs() << "    Jump is not related to this jump table\n");
@@ -241,7 +99,7 @@ MachineInstr* X86MatchJumptablePass::traceIndirectJumps(MachineFunction &MF,
             }
         }
     }
-
+    
     LLVM_DEBUG(dbgs() << "  No related indirect jump found\n");
     return nullptr;
 }
diff --git a/llvm/lib/Target/X86/X86MatchJumptablePass.h b/llvm/lib/Target/X86/X86MatchJumptablePass.h
index 779ada3d8b6d41..bec6ec9cb59729 100644
--- a/llvm/lib/Target/X86/X86MatchJumptablePass.h
+++ b/llvm/lib/Target/X86/X86MatchJumptablePass.h
@@ -20,8 +20,6 @@ namespace llvm {
 
 class X86MatchJumptablePass : public MachineFunctionPass {
 private:
-  void insertIdentifyingMarker(MachineInstr* MI, MachineFunction &MF, unsigned JTIndex);
-  void recordJumpLocation(MachineInstr* MI, MachineFunction &MF, unsigned JTIndex);
   MachineInstr* traceIndirectJumps(MachineFunction &MF, unsigned JTIndex, 
                                   MachineJumpTableInfo *JumpTableInfo);
   bool isJumpTableRelated(MachineInstr &MI, const MachineJumpTableEntry &JTEntry, 



More information about the llvm-commits mailing list