[EVM] Clear the branch condition in removeBranch for stackified code

This fixes the issue #150.

Consider MBBs that end with either:

  JUMPI
  JUMP

or

  JUMPI

When removeBranch() is called, it removes all jump instructions.
In stackified code, this disrupts the stack layout because it leaves
an unused condition value on top of the stack. To handle this, we introduced
a new codegen-only instruction called POP_COND_S, which is replaced with
POP_S in the AsmPrinter. This instruction is inserted at the end of the
MBBk after removing the conditional branch, ensuring that the unused
condition is properly cleared and the stack layout remains consistent.
In most cases, removeBranch() is immediately followed by insertBranch(),
which typically reintroduces a JUMPI instruction, possibly with a different
target. In such scenarios, the condition will again be consumed by the new
JUMPI, making POP_COND_S unnecessary. Therefore, it is removed in
insertBranch().
Although this approach introduces a minor overhead (creating and then removing
 POP_COND_S), the impact is negligible since it only applies to stackified code.
In the rare cases where removeBranch() is invoked without a subsequent
insertBranch(), the POP_COND_S instruction remains and is emitted as a regular
POP instruction.
This occurs, for example, in several places within the BranchFolder pass.

Alternative approaches considered:
  1. Identify all locations where removeBranch() is used alone and disable them
     for stackified code. This doesn't look practical and reliable.
  2: Modify analyzeBranch() to skip analyzing basic block terminators when both
     the JUMPI and JUMP (or fallthrough) targets lead to the same block. This
     would prevent the logic that calls removeBranch() from triggering.
     This approach also doesn't look reliable.

Author: PavelKopyl

llvm/lib/Target/EVM/EVMAsmPrinter.cpp

@@ -237,6 +237,12 @@ void EVMAsmPrinter::emitInstruction(const MachineInstr *MI) {
   case EVM::LOADIMMUTABLE_S:
     emitLoadImmutableLabel(MI);
     return;
+  case EVM::POP_COND_S: {
+    MCInst Pop;
+    Pop.setOpcode(EVM::POP_S);
+    EmitToStreamer(*OutStreamer, Pop);
+    return;
+  }
   }
 
   MCInst TmpInst;

llvm/lib/Target/EVM/EVMInstrInfo.cpp

@@ -173,20 +173,58 @@ unsigned EVMInstrInfo::removeBranch(MachineBasicBlock &MBB,
   LLVM_DEBUG(dbgs() << "Removing branches out of " << printMBBReference(MBB)
                     << '\n');
 
-  unsigned Count = 0;
   // Only remove branches from the end of the MBB.
-  for (auto I = MBB.rbegin(); I != MBB.rend() && I->isBranch(); ++Count) {
-
+  unsigned Count = 0;
+  bool HasConditionalBranch = false;
+  MachineBasicBlock::iterator I = MBB.end();
+  while (I != MBB.begin()) {
+    --I;
+    if (I->isDebugInstr())
+      continue;
+    if (!I->isTerminator())
+      break;
+
+    HasConditionalBranch |= I->isConditionalBranch();
 #ifndef NDEBUG
     if (I->isUnconditionalBranch())
       assert(!Count && "Malformed basic block: unconditional branch is not the "
                        "last instruction in the block");
 #endif // NDEBUG
 
     LLVM_DEBUG(dbgs() << "Removing branch: "; I->dump());
-    MBB.erase(&MBB.back());
-    I = MBB.rbegin();
+    I->eraseFromParent();
+    I = MBB.end();
+    ++Count;
   }
+
+  // Let's consider basic blocks that end with
+  //
+  //   JUMPI
+  //   JUMP
+  //
+  // or
+  //
+  //  JUMPI
+  //
+  // When removeBranch() is called, it removes all jump instructions.
+  // In stackified code, this can disrupt the stack layout by leaving an unused
+  // condition on top of the stack, a condition that would normally be consumed
+  // by the JUMPI instruction.
+  // To handle this, we add a POP_COND_S instruction at the end of the BB after
+  // removing the conditional branch, ensuring that the unused condition is
+  // properly consumed:
+  // In most cases, removeBranch() is followed by insertBranch(), which
+  // typically reinserts a JUMPI instruction, possibly with a different target.
+  // In such cases, the condition will again be consumed by the new JUMPI, so
+  // the POP_COND_S is no longer needed and is removed. Although this adds a
+  // small overhead, by creating and then removing POP_COND_S, the cost is
+  // negligible since it only affects stackified code.
+  //
+  const bool IsStackified =
+      MBB.getParent()->getInfo<EVMMachineFunctionInfo>()->getIsStackified();
+  if (IsStackified && HasConditionalBranch)
+    BuildMI(&MBB, DebugLoc(), get(EVM::POP_COND_S));
+
   return Count;
 }
 
@@ -215,6 +253,14 @@ unsigned EVMInstrInfo::insertBranch(MachineBasicBlock &MBB,
     else
       CondOpc = Cond[0].getImm() ? EVM::PseudoJUMPI : EVM::PseudoJUMP_UNLESS;
 
+    // Remove the POP_COND_S instruction previously added by removeBranch,
+    // as the newly inserted conditional branch will now consume that
+    // condition.
+    auto Last = MBB.getLastNonDebugInstr();
+    if (IsStackified && Last != MBB.end() &&
+        Last->getOpcode() == EVM::POP_COND_S)
+      MBB.erase(*Last);
+
     auto NewMI = BuildMI(&MBB, DL, get(CondOpc)).addMBB(TBB);
 
     // Add a condition operand, if we are not in stackified form.

llvm/lib/Target/EVM/EVMInstrInfo.td

@@ -830,6 +830,11 @@ let hasSideEffects = 1 in {
 
 defm POP : I<(outs), (ins), [], "POP", "", 0x50, 2>;
 
+// This instruction is a regular POP operation that indicates it clears the
+// branch condition from the top of the stack.
+let isCodeGenOnly = 1 in
+def POP_COND_S : NI<(outs), (ins), [], true, "POP", 0x50, 2>;
+
 foreach I = {1-16} in {
   defm DUP#I : I<(outs), (ins), [],
                   "DUP"#I, "", !add(I, 0x7F), 3>;

llvm/test/CodeGen/EVM/branch-folding-clear-condition.mir

@@ -0,0 +1,161 @@
+# RUN: llc -x mir -run-pass=branch-folder < %s | FileCheck %s
+# RUN: llc -x mir -start-after=evm-finalize-stack-frames < %s | FileCheck --check-prefix=ASM %s
+
+# Test that after branch folding, the condition at the top of
+# the stack is correctly cleared when a conditional branch is removed.
+
+--- |
+  ; ModuleID = 'test.ll'
+  source_filename = "test.ll"
+  target datalayout = "E-p:256:256-i256:256:256-S256-a:256:256"
+  target triple = "evm"
+
+  define void @test() {
+  entry:
+    %calldata_load_result.i21 = load i256, ptr addrspace(2) null, align 4
+    %calldata_load_result.i27 = load i256, ptr addrspace(2) inttoptr (i256 32 to ptr addrspace(2)), align 4
+    %calldata_load_result.i35 = load i256, ptr addrspace(2) inttoptr (i256 64 to ptr addrspace(2)), align 4
+    %comparison_result111.i = icmp eq i256 %calldata_load_result.i35, 0
+    br i1 %comparison_result111.i, label %if_main114.i, label %if_main215.i
+
+  if_main114.i:                                     ; preds = %entry
+    %0 = or i256 %calldata_load_result.i21, %calldata_load_result.i27
+    %or.cond119 = icmp eq i256 %0, 0
+    br i1 %or.cond119, label %if_main.i.i5, label %fun_check_entrypoint_396.exit
+
+  if_main215.i:                                     ; preds = %entry
+    %1 = or i256 %calldata_load_result.i21, %calldata_load_result.i27
+    %or.cond118 = icmp eq i256 %1, 0
+    br i1 %or.cond118, label %if_main.i.i5, label %fun_check_entrypoint_396.exit
+
+  if_main.i.i5:                                     ; preds = %if_main215.i, %if_main114.i
+     call void @llvm.evm.revert(ptr addrspace(1) null, i256 0)
+     unreachable
+
+  fun_check_entrypoint_396.exit:                    ; preds = %if_main215.i, %if_main114.i
+     call void @llvm.evm.revert(ptr addrspace(1) null, i256 32)
+     unreachable
+  }
+
+  declare void @llvm.evm.revert(ptr addrspace(1), i256)
+
+...
+# CHECK-LABEL: name: test
+# CHECK:        POP_S
+# CHECK-NEXT:   PUSH1_S i256 64
+# CHECK-NEXT:   CALLDATALOAD_S
+# CHECK-NEXT:   POP_COND_S
+# CHECK-NEXT:   PUSH1_S i256 32
+# CHECK-NEXT:   CALLDATALOAD_S
+# CHECK-NEXT:   PUSH0_S
+# CHECK-NEXT:   CALLDATALOAD_S
+# CHECK-NEXT:   OR_S
+# CHECK-NEXT:   PseudoJUMP_UNLESS %bb.2
+# CHECK-LABEL: bb.1.fun_check_entrypoint_396.exit:
+# CHECK-NEXT:   PUSH1_S i256 32
+# CHECK-NEXT:   PUSH0_S
+# CHECK-NEXT:   REVERT_S
+# CHECK-LABEL: bb.2.if_main.i.i5:
+# CHECK-NEXT:   PUSH0_S
+# CHECK-NEXT:   PUSH0_S
+# CHECK-NEXT:   REVERT_S
+
+# ASM-LABEL: test:
+# ASM:   	JUMPDEST
+# ASM-NEXT:  	POP
+# ASM-NEXT:  	PUSH1           0x40
+# ASM-NEXT:  	CALLDATALOAD
+# ASM-NEXT:  	POP
+# ASM-NEXT:  	PUSH1           0x20
+# ASM-NEXT:  	CALLDATALOAD
+# ASM-NEXT:  	PUSH0
+# ASM-NEXT:  	CALLDATALOAD
+# ASM-NEXT:  	OR
+# ASM-NEXT:  	ISZERO
+# ASM-NEXT:  	PUSH4           @.BB0_2
+# ASM-NEXT:  	JUMPI
+# ASM-LABEL: ; %bb.1:
+# ASM-NEXT:  	PUSH1           0x20
+# ASM-NEXT:  	PUSH0
+# ASM-NEXT:  	REVERT
+# ASM-LABEL: .BB0_2:
+# ASM-NEXT:  	JUMPDEST
+# ASM-NEXT:  	PUSH0
+# ASM-NEXT:  	PUSH0
+# ASM-NEXT:  	REVERT
+---
+name:            test
+alignment:       1
+machineFunctionInfo:
+  isStackified:    true
+  numberOfParameters: 0
+  hasPushDeployAddress: false
+body:             |
+  bb.0.entry:
+    successors: %bb.1(0x40000000), %bb.3(0x40000000); %bb.1(50.00%), %bb.3(50.00%)
+
+    POP_S
+    PUSH1_S i256 64
+    CALLDATALOAD_S
+    PseudoJUMPI %bb.3
+
+  bb.1.if_main114.i:
+  ; predecessors: %bb.0
+    successors: %bb.8(0x40000000), %bb.6(0x40000000); %bb.8(50.00%), %bb.6(50.00%)
+
+    PUSH1_S i256 32
+    CALLDATALOAD_S
+    PUSH0_S
+    CALLDATALOAD_S
+    OR_S
+    PseudoJUMPI %bb.6
+
+  bb.8:
+  ; predecessors: %bb.1
+    successors: %bb.4(0x80000000); %bb.4(100.00%)
+
+    PseudoJUMP %bb.4
+
+  bb.6:
+  ; predecessors: %bb.1
+    successors: %bb.2(0x80000000); %bb.2(100.00%)
+
+
+  bb.2.fun_check_entrypoint_396.exit:
+  ; predecessors: %bb.5, %bb.6
+
+    PUSH1_S i256 32
+    PUSH0_S
+    REVERT_S
+
+  bb.3.if_main215.i:
+  ; predecessors: %bb.0
+    successors: %bb.7(0x40000000), %bb.5(0x40000000); %bb.7(50.00%), %bb.5(50.00%)
+
+    PUSH1_S i256 32
+    CALLDATALOAD_S
+    PUSH0_S
+    CALLDATALOAD_S
+    OR_S
+    PseudoJUMPI %bb.5
+
+  bb.7:
+  ; predecessors: %bb.3
+    successors: %bb.4(0x80000000); %bb.4(100.00%)
+
+    PseudoJUMP %bb.4
+
+  bb.5:
+  ; predecessors: %bb.3
+    successors: %bb.2(0x80000000); %bb.2(100.00%)
+
+    PseudoJUMP %bb.2
+
+  bb.4.if_main.i.i5:
+  ; predecessors: %bb.7, %bb.8
+
+    PUSH0_S
+    PUSH0_S
+    REVERT_S
+
+...