[python-bindings] Added support for iterating over a basic blocks instructions, getting their name/dumping them, f/b iteration.

Tests are included.

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@190475 91177308-0d34-0410-b5e6-96231b3b80d8

2 files changed, 99 insertions, 0 deletions

diff --git a/bindings/python/llvm/core.py b/bindings/python/llvm/core.py
index 14b0b4ce7b1..740ce47188d 100644
--- a/bindings/python/llvm/core.py
+++ b/bindings/python/llvm/core.py
@@ -26,6 +26,7 @@ __all__ = [
     "Value",
     "Function",
     "BasicBlock",
+    "Instruction",
     "Context",
     "PassRegistry"
 ]
@@ -254,12 +255,70 @@ class BasicBlock(LLVMObject):
         return b and BasicBlock(b)
     
     @property
+    def first(self):
+        i = lib.LLVMGetFirstInstruction(self)
+        return i and Instruction(i)
+
+    @property
+    def last(self):
+        i = lib.LLVMGetLastInstruction(self)
+        return i and Instruction(i)
+
+    @property
     def name(self):
         return lib.LLVMGetValueName(Value(lib.LLVMBasicBlockAsValue(self)))
 
     def dump(self):
         lib.LLVMDumpValue(Value(lib.LLVMBasicBlockAsValue(self)))
 
+    class __inst_iterator(object):
+        def __init__(self, bb, reverse=False):            
+            self.bb = bb
+            self.reverse = reverse
+            if self.reverse:
+                self.inst = self.bb.last
+            else:
+                self.inst = self.bb.first
+        
+        def __iter__(self):
+            return self
+        
+        def next(self):
+            if not isinstance(self.inst, Instruction):
+                raise StopIteration("")
+            result = self.inst
+            if self.reverse:
+                self.inst = self.inst.prev
+            else:
+                self.inst = self.inst.next
+            return result
+    
+    def __iter__(self):
+        return BasicBlock.__inst_iterator(self)
+
+    def __reversed__(self):
+        return BasicBlock.__inst_iterator(self, reverse=True)
+
+
+class Instruction(Value):
+
+    def __init__(self, value):
+        Value.__init__(self, value)
+
+    @property
+    def next(self):
+        i = lib.LLVMGetNextInstruction(self)
+        return i and Instruction(i)
+
+    @property
+    def prev(self):
+        i = lib.LLVMGetPreviousInstruction(self)
+        return i and Instruction(i)
+
+    @property
+    def opcode(self):
+        return OpCode.from_value(lib.LLVMGetInstructionOpcode(self))
+
 class Context(LLVMObject):
 
     def __init__(self, context=None):
@@ -402,12 +461,28 @@ def register_library(library):
     library.LLVMGetPreviousBasicBlock.argtypes = [BasicBlock]
     library.LLVMGetPreviousBasicBlock.restype = c_object_p
 
+    library.LLVMGetFirstInstruction.argtypes = [BasicBlock]
+    library.LLVMGetFirstInstruction.restype = c_object_p
+
+    library.LLVMGetLastInstruction.argtypes = [BasicBlock]
+    library.LLVMGetLastInstruction.restype = c_object_p
+
     library.LLVMBasicBlockAsValue.argtypes = [BasicBlock]
     library.LLVMBasicBlockAsValue.restype = c_object_p
 
     library.LLVMCountBasicBlocks.argtypes = [Function]
     library.LLVMCountBasicBlocks.restype = c_uint
 
+    # Instruction Declarations.
+    library.LLVMGetNextInstruction.argtypes = [Instruction]
+    library.LLVMGetNextInstruction.restype = c_object_p
+
+    library.LLVMGetPreviousInstruction.argtypes = [Instruction]
+    library.LLVMGetPreviousInstruction.restype = c_object_p
+
+    library.LLVMGetInstructionOpcode.argtypes = [Instruction]
+    library.LLVMGetInstructionOpcode.restype = c_uint
+
 def register_enumerations():
     for name, value in enumerations.OpCodes:
         OpCode.register(name, value)
diff --git a/bindings/python/llvm/tests/test_core.py b/bindings/python/llvm/tests/test_core.py
index 67e294b056b..ec47c67bbc4 100644
--- a/bindings/python/llvm/tests/test_core.py
+++ b/bindings/python/llvm/tests/test_core.py
@@ -100,3 +100,27 @@ class TestCore(TestBase):
             self.assertEqual(bb.name, bb_list[i])
             bb.dump()
 
+    def test_basicblock_instruction_iteration(self):
+        m = parse_bitcode(MemoryBuffer(filename=self.get_test_bc()))
+        i = 0
+        
+        inst_list = [('arg1', OpCode.ExtractValue),
+                     ('arg2', OpCode.ExtractValue),
+                     ('', OpCode.Call),
+                     ('', OpCode.Ret)]
+        
+        bb = m.first.first
+
+        # Forward
+        for inst in bb:
+            self.assertEqual(inst.name, inst_list[i][0])
+            self.assertEqual(inst.opcode, inst_list[i][1])
+            inst.dump()
+            i += 1
+        
+        # Backwards
+        for inst in reversed(bb):
+            i -= 1
+            self.assertEqual(inst.name, inst_list[i][0])
+            self.assertEqual(inst.opcode, inst_list[i][1])
+            inst.dump()