From e6fced25b93946b75742a1f73c6a468a287c8603 Mon Sep 17 00:00:00 2001
From: Yahya Jabary <jabaryyahya@gmail.com>
Date: Tue, 30 Jun 2026 15:04:39 +0200
Subject: [PATCH] feat(ir): add OperationPtr.getOpOperands

---
 UnitTest.lean              |  1 +
 UnitTest/IR/Operation.lean | 53 ++++++++++++++++++++++++++++++++++++++
 Veir/IR/Basic.lean         | 48 ++++++++++++++++++++++++++++++++++
 3 files changed, 102 insertions(+)
 create mode 100644 UnitTest/IR/Operation.lean

diff --git a/UnitTest.lean b/UnitTest.lean
index d77fa2a27..4f97a6d67 100644
--- a/UnitTest.lean
+++ b/UnitTest.lean
@@ -3,6 +3,7 @@ import UnitTest.ParserError
 import UnitTest.Parser
 import UnitTest.AttrParser
 import UnitTest.MlirParser
+import UnitTest.IR.Operation
 import UnitTest.FP
 import UnitTest.Bitblasting.Bitblasting
 import UnitTest.DataFlowFramework.Dominance
diff --git a/UnitTest/IR/Operation.lean b/UnitTest/IR/Operation.lean
new file mode 100644
index 000000000..7d73d58f4
--- /dev/null
+++ b/UnitTest/IR/Operation.lean
@@ -0,0 +1,53 @@
+import Veir.IR.Basic
+import Veir.Parser.MlirParser
+
+open Veir
+open Veir.Parser
+
+private def parse (s : String) : OperationPtr × IRContext OpCode :=
+  match WfIRContext.create OpCode with
+  | none => panic! "failed to create IR context"
+  | some (ctx, _) =>
+    match ParserState.fromInput s.toByteArray with
+    | .error _ => panic! "lex error"
+    | .ok parser =>
+      match (parseOp none).run (MlirParserState.fromContext ctx) parser with
+      | .error _ => panic! "parse error"
+      | .ok (op, state, _) => (op, state.ctx.raw)
+
+private def flattenOps (top : OperationPtr) (ctx : IRContext OpCode) :
+    Array OperationPtr := Id.run do
+  let mut ops := #[]
+  for region in (top.get! ctx).regions do
+    let region := region.get! ctx
+    let mut currentBlock := region.firstBlock
+    while let some block := currentBlock do
+      let mut currentOp := (block.get! ctx).firstOp
+      while let some op := currentOp do
+        ops := ops.push op
+        currentOp := (op.get! ctx).next
+      currentBlock := (block.get! ctx).next
+  ops
+
+private def parsed := parse r#""builtin.module"() ({
+  %a = "test.test"() : () -> i32
+  %b = "test.test"() : () -> i32
+  %c = "arith.muli"(%a, %b) : (i32, i32) -> i32
+}) : () -> ()"#
+
+private def ops := flattenOps parsed.1 parsed.2
+
+#guard ops[0]!.getNumOperands! parsed.2 == 0
+#guard (ops[0]!.getOpOperands! parsed.2).size == 0
+
+#guard ops[2]!.getNumOperands! parsed.2 == 2
+#guard (ops[2]!.getOpOperands! parsed.2).size == ops[2]!.getNumOperands! parsed.2
+#guard (ops[2]!.getOpOperands! parsed.2)[0]! == ops[2]!.getOpOperand 0
+#guard (ops[2]!.getOpOperands! parsed.2)[1]! == ops[2]!.getOpOperand 1
+
+#guard (ops[2]!.getOpOperands! parsed.2)[0]!.op == ops[2]!
+#guard (ops[2]!.getOpOperands! parsed.2)[0]!.index == 0
+#guard (ops[2]!.getOpOperands! parsed.2)[1]!.index == 1
+
+#guard ops[2]!.getOperand! parsed.2 0 == (ops[0]!.getResult 0 : ValuePtr)
+#guard ops[2]!.getOperand! parsed.2 1 == (ops[1]!.getResult 0 : ValuePtr)
diff --git a/Veir/IR/Basic.lean b/Veir/IR/Basic.lean
index b515d4248..d6cc7fee5 100644
--- a/Veir/IR/Basic.lean
+++ b/Veir/IR/Basic.lean
@@ -490,6 +490,54 @@ theorem getOperands!.getElem_eq_getOperand! {op : OperationPtr} {h} :
     (op.getOperands! ctx)[index]'h = op.getOperand! ctx index := by
   grind [getOperands!, getOperand!]
 
+def getOpOperands (op : OperationPtr) (ctx : IRContext OpInfo)
+  (inBounds : op.InBounds ctx := by grind) : Array OpOperandPtr :=
+  Array.map (fun i => op.getOpOperand i) (Array.range (op.getNumOperands ctx inBounds))
+
+def getOpOperands! (op : OperationPtr) (ctx : IRContext OpInfo) : Array OpOperandPtr :=
+  Array.map (fun i => op.getOpOperand i) (Array.range (op.getNumOperands! ctx))
+
+@[grind =_, eq_bang ←]
+theorem getOpOperands!_eq_getOpOperands {op : OperationPtr} (hin : op.InBounds ctx) :
+    op.getOpOperands! ctx = op.getOpOperands ctx (by grind) := by
+  grind [getOpOperands, getOpOperands!]
+
+theorem getOpOperands!.mem_iff_exists_index {op : OperationPtr} :
+    operand ∈ op.getOpOperands! ctx ↔
+    ∃ index, index < op.getNumOperands! ctx ∧ op.getOpOperand index = operand := by
+  simp only [getOpOperands!, Array.mem_map, getOpOperand, getNumOperands!]
+  constructor
+  · rintro ⟨opr, ⟨hopr, oprValue⟩⟩
+    have ⟨i, hi, hopr⟩ := Array.getElem_of_mem hopr
+    exists i
+    grind
+  · grind
+
+theorem getOpOperands!.mem_getOpOperand {op : OperationPtr} :
+    index < op.getNumOperands! ctx →
+    op.getOpOperand index ∈ op.getOpOperands! ctx := by
+  grind [getOpOperands!, getOpOperand, getNumOperands!]
+
+@[simp, grind =]
+theorem getOpOperands!.size_eq_getNumOperands! {op : OperationPtr} :
+    (op.getOpOperands! ctx).size = op.getNumOperands! ctx := by
+  grind [getOpOperands!, getNumOperands!]
+
+@[simp, grind =]
+theorem getOpOperands!.getElem!_eq_getOpOperand {op : OperationPtr} :
+    index < op.getNumOperands! ctx →
+    (op.getOpOperands! ctx)[index]! = op.getOpOperand index := by
+  simp only [getOpOperands!, getOpOperand]
+  grind
+
+@[simp, grind =]
+theorem getOpOperands!.getElem_eq_getOpOperand
+    {op : OperationPtr} {h : index < (op.getOpOperands! ctx).size} :
+    index < op.getNumOperands! ctx →
+    (op.getOpOperands! ctx)[index]'h = op.getOpOperand index := by
+  simp only [getOpOperands!, getOpOperand]
+  grind
+
 def getOperandTypes (op : OperationPtr) (ctx : IRContext OpInfo)
     (inBounds : op.InBounds ctx := by grind) : Array TypeAttr :=
   (op.get ctx).operands.map fun opr =>