[python-bindings] Added support for iterating over a function's basic blocks, dumping/getting names of those bb, f/w iteration.

Tests are included.

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

2 files changed, 105 insertions, 0 deletions

diff --git a/bindings/python/llvm/core.py b/bindings/python/llvm/core.py
index 3da69d3..14b0b4c 100644
--- a/bindings/python/llvm/core.py
+++ b/bindings/python/llvm/core.py
@@ -25,6 +25,7 @@ __all__ = [
     "Module",
     "Value",
     "Function",
+    "BasicBlock",
     "Context",
     "PassRegistry"
 ]
@@ -196,6 +197,69 @@ class Function(Value):
         f = lib.LLVMGetPreviousFunction(self)
         return f and Function(f)
     
+    @property
+    def first(self):
+        b = lib.LLVMGetFirstBasicBlock(self)
+        return b and BasicBlock(b)
+
+    @property
+    def last(self):
+        b = lib.LLVMGetLastBasicBlock(self)
+        return b and BasicBlock(b)
+
+    class __bb_iterator(object):
+        def __init__(self, function, reverse=False):
+            self.function = function
+            self.reverse = reverse
+            if self.reverse:
+                self.bb = function.last
+            else:
+                self.bb = function.first
+        
+        def __iter__(self):
+            return self
+        
+        def next(self):
+            if not isinstance(self.bb, BasicBlock):
+                raise StopIteration("")
+            result = self.bb
+            if self.reverse:
+                self.bb = self.bb.prev
+            else:
+                self.bb = self.bb.next
+            return result
+    
+    def __iter__(self):
+        return Function.__bb_iterator(self)
+
+    def __reversed__(self):
+        return Function.__bb_iterator(self, reverse=True)
+    
+    def __len__(self):
+        return lib.LLVMCountBasicBlocks(self)
+
+class BasicBlock(LLVMObject):
+    
+    def __init__(self, value):
+        LLVMObject.__init__(self, value)
+
+    @property
+    def next(self):
+        b = lib.LLVMGetNextBasicBlock(self)
+        return b and BasicBlock(b)
+
+    @property
+    def prev(self):
+        b = lib.LLVMGetPreviousBasicBlock(self)
+        return b and BasicBlock(b)
+    
+    @property
+    def name(self):
+        return lib.LLVMGetValueName(Value(lib.LLVMBasicBlockAsValue(self)))
+
+    def dump(self):
+        lib.LLVMDumpValue(Value(lib.LLVMBasicBlockAsValue(self)))
+
 class Context(LLVMObject):
 
     def __init__(self, context=None):
@@ -325,6 +389,25 @@ def register_library(library):
     library.LLVMDumpValue.argtypes = [Value]
     library.LLVMDumpValue.restype = None
 
+    # Basic Block Declarations.
+    library.LLVMGetFirstBasicBlock.argtypes = [Function]
+    library.LLVMGetFirstBasicBlock.restype = c_object_p
+
+    library.LLVMGetLastBasicBlock.argtypes = [Function]
+    library.LLVMGetLastBasicBlock.restype = c_object_p
+
+    library.LLVMGetNextBasicBlock.argtypes = [BasicBlock]
+    library.LLVMGetNextBasicBlock.restype = c_object_p
+
+    library.LLVMGetPreviousBasicBlock.argtypes = [BasicBlock]
+    library.LLVMGetPreviousBasicBlock.restype = c_object_p
+
+    library.LLVMBasicBlockAsValue.argtypes = [BasicBlock]
+    library.LLVMBasicBlockAsValue.restype = c_object_p
+
+    library.LLVMCountBasicBlocks.argtypes = [Function]
+    library.LLVMCountBasicBlocks.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 a1f79a4..67e294b 100644
--- a/bindings/python/llvm/tests/test_core.py
+++ b/bindings/python/llvm/tests/test_core.py
@@ -78,3 +78,25 @@ class TestCore(TestBase):
             self.assertEqual(f.name, functions[i])
             f.dump()
 
+    def test_function_basicblock_iteration(self):
+        m = parse_bitcode(MemoryBuffer(filename=self.get_test_bc()))
+        i = 0
+        
+        bb_list = ['b1', 'b2', 'end']
+        
+        f = m.first
+        while f.name != "f6":
+            f = f.next
+        
+        # Forward
+        for bb in f:
+            self.assertEqual(bb.name, bb_list[i])
+            bb.dump()
+            i += 1
+        
+        # Backwards
+        for bb in reversed(f):
+            i -= 1
+            self.assertEqual(bb.name, bb_list[i])
+            bb.dump()
+