Bugfix: call partition in basic blocks

Calls must be partitioned like jp/jr. The "next" block is not
the one upon return.

Also adds some testing for basic blocks

Author: boriel

arch/zx48k/optimizer/basicblock.py

@@ -332,6 +332,10 @@ def update_goes_and_comes(self):
                 self.add_goes_to(self.next)
 
         elif inst == 'call':
+            if cond is None:
+                self.delete_goes(self.next)
+                self.next = LABELS[oper[0]].basic_block
+
             LABELS[oper[0]].basic_block.add_comes_from(self)
             stack = [LABELS[oper[0]].basic_block]
             bbset = IdentitySet()

tests/arch/zx48k/optimizer/test_basicblock.py

@@ -0,0 +1,67 @@
+# -*- coding: utf-8 -*-
+
+import unittest
+from arch.zx48k.optimizer import basicblock
+
+
+class TestBasicBlock(unittest.TestCase):
+    """ Tests basic blocks implementation
+    """
+    def setUp(self):
+        self.blk = basicblock.BasicBlock([])
+
+    def test_block_partition_jp(self):
+        code = """
+        nop
+        jp __UNKNOWN
+        nop
+        """
+        self.blk.code = [x for x in code.split('\n') if x.strip()]
+        blks = basicblock.get_basic_blocks(self.blk)
+        self.assertEqual(len(blks), 2)
+        self.assertEqual(blks[0].code, ['nop', 'jp __UNKNOWN'])
+        self.assertEqual(blks[1].code, ['nop'])
+        self.assertFalse(blks[1] in blks[0].goes_to)
+        self.assertFalse(blks[0] in blks[1].comes_from)
+
+    def test_block_partition_call(self):
+        code = """
+        nop
+        call __UNKNOWN
+        nop
+        """
+        self.blk.code = [x for x in code.split('\n') if x.strip()]
+        blks = basicblock.get_basic_blocks(self.blk)
+        self.assertEqual(len(blks), 2)
+        self.assertEqual(blks[0].code, ['nop', 'call __UNKNOWN'])
+        self.assertEqual(blks[1].code, ['nop'])
+        self.assertFalse(blks[1] in blks[0].goes_to)
+        self.assertFalse(blks[0] in blks[1].comes_from)
+
+    def test_block_partition_jp_flag(self):
+        code = """
+        nop
+        jp z, __UNKNOWN
+        nop
+        """
+        self.blk.code = [x for x in code.split('\n') if x.strip()]
+        blks = basicblock.get_basic_blocks(self.blk)
+        self.assertEqual(len(blks), 2)
+        self.assertEqual(blks[0].code, ['nop', 'jp z, __UNKNOWN'])
+        self.assertEqual(blks[1].code, ['nop'])
+        self.assertTrue(blks[1] in blks[0].goes_to)
+        self.assertTrue(blks[0] in blks[1].comes_from)
+
+    def test_block_partition_call_flag(self):
+        code = """
+        nop
+        call z, __UNKNOWN
+        nop
+        """
+        self.blk.code = [x for x in code.split('\n') if x.strip()]
+        blks = basicblock.get_basic_blocks(self.blk)
+        self.assertEqual(len(blks), 2)
+        self.assertEqual(blks[0].code, ['nop', 'call z, __UNKNOWN'])
+        self.assertEqual(blks[1].code, ['nop'])
+        self.assertTrue(blks[1] in blks[0].goes_to)
+        self.assertTrue(blks[0] in blks[1].comes_from)