fix pack, survey and move to doctest

2 files changed, 146 insertions, 148 deletions

diff --git a/pack.py b/pack.py
index 98d7d86..f57d645 100755
--- a/pack.py
+++ b/pack.py
@@ -1,11 +1,60 @@
 #!/usr/bin/env python3
+""" pack.py
+Packs bitfields tagged with their size into untagged bytes.
+
+>>> from io import StringIO
+>>> # doctest: +REPORT_NDIFF
+... print(subv.join_all(pack(StringIO('''
+... == code 0x80000000
+... 37/7 05/5 10010/20
+... 13/7 06/5 00/3 00/5 48/12
+... 23/7 00/5 02/3 05/5 06/5 00/7
+... 13/7 06/5 00/3 00/5 65/12
+... 23/7 00/5 02/3 05/5 06/5 00/7
+... 13/7 06/5 00/3 00/5 6c/12
+... 23/7 00/5 02/3 05/5 06/5 00/7
+... 13/7 06/5 00/3 00/5 6c/12
+... 23/7 00/5 02/3 05/5 06/5 00/7
+... 13/7 06/5 00/3 00/5 6f/12
+... 23/7 00/5 02/3 05/5 06/5 00/7
+... 13/7 06/5 00/3 00/5 0a/12
+... 23/7 00/5 02/3 05/5 06/5 00/7
+... 6f/7 00/5 ff/8 01/1 3e6/10 01/1
+... '''[1:-1]))))
+== code 0x80000000
+b7 02 01 10
+13 03 80 04
+23 a0 62 00
+13 03 50 06
+23 a0 62 00
+13 03 c0 06
+23 a0 62 00
+13 03 c0 06
+23 a0 62 00
+13 03 f0 06
+23 a0 62 00
+13 03 a0 00
+23 a0 62 00
+6f f0 df fc
+"""
+
 import subv
 import bits
 
 def byteify(word):
+    """ split longer bitfield into bytes.
+
+    >>> byteify((0x12345678, 32))
+    [(120,), (86,), (52,), (18,)]
+    >>> byteify((0x4801813, 32))
+    [(19,), (24,), (128,), (4,)]
+    >>> byteify((0x13, 9)),
+    Traceback (most recent call last):
+        ...
+    AssertionError: not byte aligned: 9 bits
+    """
     (val, size) = word
-    if size % 8 != 0:
-        raise ValueError("Not byte aligned: {} bits".format(size))
+    assert size % 8 == 0, "not byte aligned: {} bits".format(size)
 
     out = []
     for i in range(0, size, 8):
@@ -21,10 +70,8 @@ def pack(iter):
             fields = [bits.from_part(p) for p in line['instr']]
             total = bits.concat(*fields)
 
-            if segment == 'code' and total[1] != 32:
-                raise ValueError("instruction parts do not add up to 32 bit!")
-
-            yield subv.format_instr(byteify(total))
+            line['instr'] =  byteify(total)
+            yield subv.format(line)
         else:
             if line['type'] == 'segment':
                 segment = line['segment'][0]
@@ -34,62 +81,3 @@ if __name__ == '__main__':
     import sys
     for line in pack(sys.stdin):
         print(line)
-
-import unittest
-class TestPack(unittest.TestCase):
-    def test_byteify(self):
-        self.assertEqual(
-            byteify((0x12345678, 32)),
-            [(0x78,), (0x56,), (0x34,), (0x12,)]
-        )
-        self.assertEqual(
-            byteify((0x4801813, 32)),
-            [(0x13,), (0x18,), (0x80,), (0x04,)]
-        )
-
-    def test_e2e(self):
-        from io import StringIO
-        from textwrap import dedent
-        inp = dedent('''\
-            # do some things
-            == code 0x80000000
-                37/7 05/5 10010/20
-                13/7 06/5 00/3 00/5 48/12
-                23/7 00/5 02/3 05/5 06/5 00/7
-                13/7 06/5 00/3 00/5 65/12
-                23/7 00/5 02/3 05/5 06/5 00/7
-                13/7 06/5 00/3 00/5 6c/12
-                23/7 00/5 02/3 05/5 06/5 00/7
-                13/7 06/5 00/3 00/5 6c/12
-                23/7 00/5 02/3 05/5 06/5 00/7
-                13/7 06/5 00/3 00/5 6f/12
-                23/7 00/5 02/3 05/5 06/5 00/7
-                13/7 06/5 00/3 00/5 0a/12
-                23/7 00/5 02/3 05/5 06/5 00/7
-                6f/7 00/5 ff/8 01/1 3e6/10 1/1
-        ''')
-
-        out = dedent('''\
-            # do some things
-            == code 0x80000000
-            b7 02 01 10
-            13 03 80 04
-            23 a0 62 00
-            13 03 50 06
-            23 a0 62 00
-            13 03 c0 06
-            23 a0 62 00
-            13 03 c0 06
-            23 a0 62 00
-            13 03 f0 06
-            23 a0 62 00
-            13 03 a0 00
-            23 a0 62 00
-            6f f0 df fc
-        ''')
-
-        got = ''
-        for line in pack(StringIO(inp)):
-            got += line + '\n'
-
-        self.assertEqual(got, out)
diff --git a/survey.py b/survey.py
index e3d7a3e..bf9e551 100755
--- a/survey.py
+++ b/survey.py
@@ -1,4 +1,51 @@
 #!/usr/bin/env python3
+""" survey.py
+Resolves label references and substitutes them with literal values.
+
+>>> from io import StringIO
+>>> # doctest: +REPORT_NDIFF
+... print(subv.join_all(survey(StringIO('''
+... == code 0x80000000
+... main:
+... # load 0x10010000 (UART0) into t0
+... 37/7 05/5 10010/20
+... # store 0x48 (H) in UART0+0
+... 13/7 06/5 00/3 00/5 48/12
+... 23/7 00/5 02/3 05/5 06/5 00/7
+... # store 0x65 (e) in UART0+0
+... 13/7 06/5 00/3 00/5 65/12
+... 23/7 00/5 02/3 05/5 06/5 00/7
+... # store 0x6c (l) in UART0+0
+... 13/7 06/5 00/3 00/5 6c/12
+... 23/7 00/5 02/3 05/5 06/5 00/7
+... # store 0x6c (l) in UART0+0
+... 13/7 06/5 00/3 00/5 6c/12
+... 23/7 00/5 02/3 05/5 06/5 00/7
+... # store 0x6f (o) in UART0+0
+... 13/7 06/5 00/3 00/5 6f/12
+... 23/7 00/5 02/3 05/5 06/5 00/7
+... # store 0x0a (\\\\n) in UART0+0
+... 13/7 06/5 00/3 00/5 0a/12
+... 23/7 00/5 02/3 05/5 06/5 00/7
+... # jump back up to the top
+... 6f/7 00/5 main[19:12]/off8 main[11:11]/off1 main[10:1]/off10 main[20:20]/off1
+... '''[1:-1]))))
+== code 0x80000000
+37/7 05/5 10010/20
+13/7 06/5 00/3 00/5 48/12
+23/7 00/5 02/3 05/5 06/5 00/7
+13/7 06/5 00/3 00/5 65/12
+23/7 00/5 02/3 05/5 06/5 00/7
+13/7 06/5 00/3 00/5 6c/12
+23/7 00/5 02/3 05/5 06/5 00/7
+13/7 06/5 00/3 00/5 6c/12
+23/7 00/5 02/3 05/5 06/5 00/7
+13/7 06/5 00/3 00/5 6f/12
+23/7 00/5 02/3 05/5 06/5 00/7
+13/7 06/5 00/3 00/5 0a/12
+23/7 00/5 02/3 05/5 06/5 00/7
+6f/7 00/5 ff/8 01/1 3e6/10 01/1
+"""
 import subv
 import bits
 import re
@@ -6,39 +53,46 @@ import re
 slice_re = re.compile(r'^([^\[]+)(?:\[(\d+):(\d+)\])?$')
 field_re = re.compile(r'^(imm|off)(\d+)$')
 
-def scan(iter):
-    addr = None
-    map = {}
-    for line in iter:
-        if line['type'] == 'label':
-            map[line['label']] = line['addr']
-    return map
-
 def observe(part, rel_addr, map):
-    if subv.is_reference(part):
-        label, trim, tag = part
-        trim = int(trim)
-
-        if label not in map:
-            raise ValueError("label not found: {} ({})".format(label, part))
-        addr = map[label]
-
-        mode, size, shift = tag_re.match(tag).groups()
-        size = int(size)
-        shift = int(shift) if shift else 0
-
-        if mode == 'imm':
-            addr = bits.u(addr, size)
-        else:
-            addr = bits.i(addr - rel_addr, size)
-
-        part = bits.slice(addr, trim+shift-1, shift)
-
-    return part
+    """ resolve a label reference.
+
+    >>> observe((3, '3'), 0, {})
+    (3, '3')
+    >>> observe(('label[31:0]', 'imm32'), 0, {'label': 1234})
+    (1234, 32)
+    >>> observe(('label[1:0]', 'imm32'), 0, {'label': 1234})
+    (2, 2)
+    >>> observe(('label[31:0]', 'off32'), 1000, {'label': 1234})
+    (234, 32)
+    >>> observe(('label[31:12]', 'off32'), 1000, {'label': 1234})
+    (0, 20)
+    >>> observe(('label[31:0]', 'off32'), 2000, {'label': 1000})
+    (4294966296, 32)
+
+    >>> observe(('label[32:0]', 'imm32'), 0, {})
+    Traceback (most recent call last):
+        ...
+    AssertionError: undefined label 'label'
+    """
+    if not subv.is_reference(part):
+        return part
+
+    ref = subv.parse_reference(part)
+    assert ref['label'] in map, "undefined label '{}'".format(ref['label'])
+    addr = map[ref['label']]
+
+    # @TODO: the hardcoded 32 here is not right, this is going to blow up
+    # in some circumstances (e.g. a backwards branch 2<<18 bytes away)
+    if ref['mode'] == 'imm':
+        addr = bits.u(addr, 32)
+    else:
+        addr = bits.i(addr - rel_addr, 32)
+
+    return bits.slice(addr, ref['hi'], ref['lo'])
 
 def survey(iter):
-    # buffer in everything first, we need two passes
     queue = []
+    map = {}
     addr = -1
     for line in iter:
         line = subv.parse(line)
@@ -48,24 +102,28 @@ def survey(iter):
         # step forward addr
         if line['type'] == 'segment':
             segment, addr = line['segment']
+        elif line['type'] == 'label':
+            map[line['label']] = addr
         elif line['type'] == 'instr':
             bits = 0
             for part in line['instr']:
-                bits += int(part[1])
-
-            if bits % 8 != 0:
-                raise ValueError("SubV lines should be byte aligned!")
+                if subv.is_reference(part):
+                  ref = subv.parse_reference(part)
+                  bits += ref['size']
+                else:
+                  bits += int(part[1])
 
+            assert bits % 8 == 0, "line not byte-aligned"
             addr += bits // 8
 
-    map = scan(queue)
     for line in queue:
         if line['type'] == 'instr':
             instr = []
             for part in line['instr']:
                 observed = observe(part, line['addr'], map)
                 instr.append(observed)
-            yield subv.format_instr(instr)
+            line['instr'] = instr
+            yield subv.format(line)
         elif line['type'] == 'segment':
             yield line['raw']
 
@@ -73,51 +131,3 @@ if __name__ == '__main__':
     import sys
     for line in survey(sys.stdin):
         print(line)
-
-import unittest
-class TestE2E(unittest.TestCase):
-    def test_e2e(self):
-        from io import StringIO
-        from textwrap import dedent
-        inp = dedent('''\
-            == code 0x80000000
-            main:
-            37/7 05/5 10010/20
-            13/7 06/5 00/3 00/5 48/12
-            23/7 00/5 02/3 05/5 06/5 00/7
-            13/7 06/5 00/3 00/5 65/12
-            23/7 00/5 02/3 05/5 06/5 00/7
-            13/7 06/5 00/3 00/5 6c/12
-            23/7 00/5 02/3 05/5 06/5 00/7
-            13/7 06/5 00/3 00/5 6c/12
-            23/7 00/5 02/3 05/5 06/5 00/7
-            13/7 06/5 00/3 00/5 6f/12
-            23/7 00/5 02/3 05/5 06/5 00/7
-            13/7 06/5 00/3 00/5 0a/12
-            23/7 00/5 02/3 05/5 06/5 00/7
-            6f/7 00/5 main/8/off21>>12 main/1/off21>>11 main/10/off21>>1 main/1/off21>>20
-        ''')
-
-        out = dedent('''\
-            == code 0x80000000
-            37/7 05/5 10010/20
-            13/7 06/5 00/3 00/5 48/12
-            23/7 00/5 02/3 05/5 06/5 00/7
-            13/7 06/5 00/3 00/5 65/12
-            23/7 00/5 02/3 05/5 06/5 00/7
-            13/7 06/5 00/3 00/5 6c/12
-            23/7 00/5 02/3 05/5 06/5 00/7
-            13/7 06/5 00/3 00/5 6c/12
-            23/7 00/5 02/3 05/5 06/5 00/7
-            13/7 06/5 00/3 00/5 6f/12
-            23/7 00/5 02/3 05/5 06/5 00/7
-            13/7 06/5 00/3 00/5 0a/12
-            23/7 00/5 02/3 05/5 06/5 00/7
-            6f/7 00/5 ff/8 01/1 3e6/10 01/1
-        ''')
-
-        got = ''
-        for line in survey(StringIO(inp)):
-            got += line + '\n'
-
-        self.assertEqual(got, out)