feat(interfaces): getSuccessorForOperands in ControlFlowInterface

Author: axelcool1234

Veir/Interfaces/ControlFlowInterfaces.lean

@@ -5,8 +5,9 @@ public import Veir.GlobalOpInfo
 /-!
 # ControlFlowInterfaces
 
-This file provides support for querying which operands are forwarded to a successor
-and mapping those operands to successor block arguments.
+This file provides support for querying which operands are forwarded to a successor,
+mapping those operands to successor block arguments, and selecting a successor from
+known constant operands.
 -/
 
 namespace Veir
@@ -29,6 +30,18 @@ instance : GetElem? SuccessorOperands Nat ValuePtr
 
 namespace BranchOpInterface
 
+/-- Return the known value of an operand, or `none` when it is unknown. -/
+private def getKnownOperand?
+    (operands : Array (Option RuntimeValue)) (index : Nat) : Option RuntimeValue := do
+  let operand ← operands[index]?
+  operand
+
+/-- Return the true or false successor of a conditional branch. -/
+private def getConditionalSuccessor?
+    (branchOp : OperationPtr) (condition : Bool) (raw : IRContext OpCode) : Option BlockPtr :=
+  let successors := branchOp.getSuccessors! raw
+  if condition then successors[0]? else successors[1]?
+
 /--
   Return the operands passed to `successorIndex` of a branch operation.
 -/
@@ -82,6 +95,52 @@ def getSuccessorOperand?
   getSuccessorOperands? branchOp successorIndex raw >>= fun operands =>
     operands[blockArgumentIndex]?
 
+/--
+Return the successor selected by the known constant operands of a branch operation.
+An operand is `none` when its value is unknown. Returns `none` when the operation is
+not a supported branch or a single successor cannot be determined.
+-/
+def getSuccessorForOperands?
+    (branchOp : OperationPtr) (operands : Array (Option RuntimeValue))
+    (raw : IRContext OpCode) : Option BlockPtr :=
+  match branchOp.getOpType! raw with
+  | .cf .br | .llvm .br | .riscv_cf .branch =>
+    (branchOp.getSuccessors! raw)[0]?
+  | .cf .cond_br | .llvm .cond_br => do
+    let .int _ (.val condition) ← getKnownOperand? operands 0 | none
+    getConditionalSuccessor? branchOp (condition ≠ 0) raw
+  | .riscv_cf .beqz => do
+    let .reg condition ← getKnownOperand? operands 0 | none
+    getConditionalSuccessor? branchOp (condition.val = 0#64) raw
+  | .riscv_cf .bnez => do
+    let .reg condition ← getKnownOperand? operands 0 | none
+    getConditionalSuccessor? branchOp (condition.val ≠ 0#64) raw
+  | .riscv_cf .beq => do
+    let .reg lhs ← getKnownOperand? operands 0 | none
+    let .reg rhs ← getKnownOperand? operands 1 | none
+    getConditionalSuccessor? branchOp (lhs = rhs) raw
+  | .riscv_cf .bne => do
+    let .reg lhs ← getKnownOperand? operands 0 | none
+    let .reg rhs ← getKnownOperand? operands 1 | none
+    getConditionalSuccessor? branchOp (lhs ≠ rhs) raw
+  | .riscv_cf .blt => do
+    let .reg lhs ← getKnownOperand? operands 0 | none
+    let .reg rhs ← getKnownOperand? operands 1 | none
+    getConditionalSuccessor? branchOp (BitVec.slt lhs.val rhs.val) raw
+  | .riscv_cf .bge => do
+    let .reg lhs ← getKnownOperand? operands 0 | none
+    let .reg rhs ← getKnownOperand? operands 1 | none
+    getConditionalSuccessor? branchOp (!BitVec.slt lhs.val rhs.val) raw
+  | .riscv_cf .bltu => do
+    let .reg lhs ← getKnownOperand? operands 0 | none
+    let .reg rhs ← getKnownOperand? operands 1 | none
+    getConditionalSuccessor? branchOp (BitVec.ult lhs.val rhs.val) raw
+  | .riscv_cf .bgeu => do
+    let .reg lhs ← getKnownOperand? operands 0 | none
+    let .reg rhs ← getKnownOperand? operands 1 | none
+    getConditionalSuccessor? branchOp (!BitVec.ult lhs.val rhs.val) raw
+  | _ => none
+
 end BranchOpInterface
 
 end