[EVM] Fold Or Jumpi

Author: akiramenai

llvm/lib/Target/EVM/EVMLowerJumpUnless.cpp

@@ -82,10 +82,16 @@ static void lowerJumpUnless(MachineInstr &MI, const EVMInstrInfo *TII,
 /// Returns `true` if any fold was performed.
 static bool tryFoldJumpUnless(MachineInstr &MI, const EVMInstrInfo *TII) {
   auto I = MachineBasicBlock::iterator(&MI);
-  auto *PrevMI = I == MI.getParent()->begin() ? nullptr : &*std::prev(I);
-  bool CanFold = PrevMI && (PrevMI->getOpcode() == EVM::ISZERO_S ||
-                            PrevMI->getOpcode() == EVM::EQ_S ||
-                            PrevMI->getOpcode() == EVM::SUB_S);
+  if (I == MI.getParent()->begin())
+    return false;
+  auto PrevMI = std::prev(I);
+  while (PrevMI != MI.getParent()->begin() && PrevMI->getOpcode() == EVM::OR_S)
+    PrevMI = std::prev(PrevMI);
+  if (PrevMI == MI.getParent()->begin())
+    return false;
+  bool CanFold = PrevMI->getOpcode() == EVM::ISZERO_S ||
+                 PrevMI->getOpcode() == EVM::EQ_S ||
+                 PrevMI->getOpcode() == EVM::SUB_S;
 
   if (!CanFold)
     return false;
@@ -121,6 +127,48 @@ static void lowerPseudoJumpUnless(MachineInstr &MI, const EVMInstrInfo *TII,
       .add(MI.getOperand(0));
 }
 
+static void optimizePseudoJumpI(MachineInstr &MI, const EVMInstrInfo *TII,
+                                const bool IsStackified,
+                                const bool FoldJumps) {
+  assert(IsStackified && "Found pseudo jump_unless in non-stackified function");
+  assert(MI.getNumExplicitOperands() == 1 &&
+         "Unexpected number of operands in pseudo jump_unless");
+
+  if (!FoldJumps)
+    return;
+  auto I = MachineBasicBlock::iterator(&MI);
+  if (I == MI.getParent()->begin())
+    return;
+  auto PrevMI = std::prev(I);
+  while (PrevMI != MI.getParent()->begin() && PrevMI->getOpcode() == EVM::OR_S)
+    PrevMI = std::prev(PrevMI);
+  if (PrevMI == MI.getParent()->begin() ||  PrevMI->getOpcode() != EVM::ISZERO_S)
+    return;
+  PrevMI = std::prev(PrevMI);
+  if (PrevMI == MI.getParent()->begin())
+    return;
+  bool CanFold = PrevMI->getOpcode() == EVM::ISZERO_S ||
+                 PrevMI->getOpcode() == EVM::EQ_S ||
+                 PrevMI->getOpcode() == EVM::SUB_S;
+
+  if (!CanFold)
+    return;
+
+  ++NumPseudoJumpUnlessFolded;
+
+  if (PrevMI->getOpcode() == EVM::ISZERO_S)
+    PrevMI->eraseFromParent();
+  else if (PrevMI->getOpcode() == EVM::EQ_S) {
+    if (auto PrevPrevMI = std::prev(PrevMI); PrevPrevMI != MI.getParent()->begin() && PrevPrevMI->getOpcode() == EVM::PUSH0_S) {
+      PrevPrevMI->eraseFromParent();
+      PrevMI->eraseFromParent();
+    } else
+      PrevMI->setDesc(TII->get(EVM::SUB_S));
+  }
+  else if (PrevMI->getOpcode() == EVM::SUB_S)
+    PrevMI->setDesc(TII->get(EVM::EQ_S));
+}
+
 bool EVMLowerJumpUnless::runOnMachineFunction(MachineFunction &MF) {
   LLVM_DEBUG({
     dbgs() << "********** Lower jump_unless instructions **********\n"
@@ -147,6 +195,9 @@ bool EVMLowerJumpUnless::runOnMachineFunction(MachineFunction &MF) {
     case EVM::JUMP_UNLESS:
       lowerJumpUnless(*TermIt, TII, IsStackified, MRI);
       break;
+    case EVM::PseudoJUMPI:
+      optimizePseudoJumpI(*TermIt, TII, IsStackified, OptLevel != CodeGenOptLevel::None);
+      break;
     default:
       continue;
     }