Improved comment alignment (#13)

4 files changed, 317 insertions, 46 deletions

diff --git a/README.md b/README.md
index 8a40780..71cc51f 100644
--- a/README.md
+++ b/README.md
@@ -39,6 +39,13 @@ With `-s` the tool adds a subtotal.
 
 By default `z80count` will try to update comments replacing existing annotations.
 
+Comments added by `z80count` are aligned to the column given in the
+`-c` (`--column`) option (50 by default). By default the comments are
+aligned using spaces, if you prefer tabs instead use the `--use-tabs`
+option. In order to compute the padding `z80count` assumes that a
+`TAB` equals 4 spaces. Use the option `-t` to override this.
+
+
 Example:
 ```asm
 	push hl
@@ -91,7 +98,7 @@ Processed with `z80count.py -s` results in:
 ```asm
 	push hl				; [11 .. 11]
 	pop bc				; [10 .. 21]
-	ld hl, $5800			; [10 .. 31]
+	ld hl, $5800		; [10 .. 31]
 
 	ld e, 7				; [7 .. 38]
 .fade_out_all_loop0
@@ -102,7 +109,7 @@ Processed with `z80count.py -s` results in:
 .fade_out_all_loop1
 	ld a, (hl)			; [7 .. 71]
 	and 7				; [7 .. 78]
-	jr z, no_fade_all_ink		; [12/7 .. 90/85]
+	jr z, no_fade_all_ink ; [12/7 .. 90/85]
 	dec a				; [4 .. 89]
 .no_fade_all_ink
 
@@ -110,7 +117,7 @@ Processed with `z80count.py -s` results in:
 
 	ld a, (hl)			; [7 .. 100]
 	and $38				; [7 .. 107]
-	jr z, no_fade_all_paper		; [12/7 .. 119/114]
+	jr z, no_fade_all_paper	; [12/7 .. 119/114]
 	sub 8				; [7 .. 121]
 .no_fade_all_paper
 
@@ -127,12 +134,12 @@ Processed with `z80count.py -s` results in:
 	dec bc				; [6 .. 166]
 	ld a, b				; [4 .. 170]
 	or c				; [4 .. 174]
-	jr nz, fade_out_all_loop1	; [12/7 .. 186/181]
+	jr nz, fade_out_all_loop1 ; [12/7 .. 186/181]
 
 	pop bc				; [10 .. 191]
 	pop hl				; [10 .. 201]
 	dec e				; [4 .. 205]
-	jr nz, fade_out_all_loop0	; [12/7 .. 217/212]
+	jr nz, fade_out_all_loop0 ; [12/7 .. 217/212]
 ```
 
 Comments show subtotals, and there are two types:
diff --git a/tests/test_format_line.py b/tests/test_format_line.py
new file mode 100644
index 0000000..d73b25e
--- /dev/null
+++ b/tests/test_format_line.py
@@ -0,0 +1,151 @@
+# -*- coding: utf-8 -*-
+
+import pytest
+
+from z80count.z80count import comment_alignment
+from z80count.z80count import format_line
+from z80count.z80count import line_length
+
+
+##########################################################################
+# test support funtions                                                  #
+##########################################################################
+
+@pytest.mark.parametrize("line,tab_width,expected", (
+    ("",      8, 0),
+    (" ",     8, 1),
+    ("a",     8, 1),
+    ("á",     8, 1),
+
+    ("\t",    8, 8),
+    (" \t",   8, 8),
+    ("  \t",  8, 8),
+    ("\t ",   8, 9),
+    ("\t\t",  8, 16),
+    ("\t \t", 8, 16),
+
+    ("\t",    3, 3),
+    (" \t",   3, 3),
+    ("  \t",  3, 3),
+    ("\t ",   3, 4),
+    ("\t\t",  3, 6),
+    ("\t \t", 3, 6),
+
+    ("\tadd hl,de", 4, 13),
+))
+def test_line_length(line, tab_width, expected):
+    assert line_length(line, tab_width) == expected
+
+
+@pytest.mark.parametrize("line,column,expected", (
+    ("foo",      8, "    "),
+    (" foo",     8, "   "),
+    ("foo  ",    8, "  "),
+    ("foo\t",    8, " "),
+    ("foo     ", 8, " "),
+
+    ("foo",      9, "\t"),
+    (" foo",     9, "\t"),
+    ("foo  ",    9, "\t"),
+    ("foo\t",    9, " "),
+    ("foo     ", 9, " "),
+
+    ("longer than tab stop", 9, " "),
+))
+def test_comment_alignment_with_tabs(line, column, expected):
+    assert comment_alignment(line, column, use_tabs=True, tab_width=8) == expected
+
+
+def test_comment_alignment_bug_001():
+    line = "\tld hl,PLY_Interruption_Convert"
+    out = comment_alignment(line, column=50, use_tabs=True, tab_width=4)
+    assert out == "\t\t\t\t"
+
+
+def test_comment_alignment_bug_002():
+    line = "\tadd hl,de"
+    out = comment_alignment(line, column=50, use_tabs=True, tab_width=4)
+    assert out == "\t\t\t\t\t\t\t\t\t"
+
+
+def test_comment_alignment_bug_003():
+    line = "\tjp #bce0"
+    out = comment_alignment(line, column=50, use_tabs=True, tab_width=4)
+    assert out == "\t\t\t\t\t\t\t\t\t"
+
+
+@pytest.mark.parametrize("line,column,expected", (
+    ("foo",      8, "    "),
+    (" foo",     8, "   "),
+    ("foo  ",    8, "  "),
+    ("foo\t",    8, " "),
+    ("foo    ",  8, " "),
+    ("foo     ", 8, " "),
+
+    ("foo",       9, "     "),
+    (" foo",      9, "    "),
+    ("foo  ",     9, "   "),
+    ("foo\t",     9, " "),
+    ("foo    ",   9, " "),
+    ("foo     ",  9, " "),
+    ("foo      ", 9, " "),
+
+    ("longer than tab stop", 9, " "),
+))
+def test_comment_alignment_with_spaces(line, column, expected):
+    assert comment_alignment(line, column, use_tabs=False, tab_width=8) == expected
+
+
+##########################################################################
+# test format_line                                                       #
+##########################################################################
+
+
+def _make_entry(cycles, case=""):
+    return {
+        "cycles": cycles,
+        "case": case,
+    }
+
+
+def _run(line, cycles, total=0, total_cond=0, case="", subt=False,
+         update=False, column=15, debug=False, use_tabs=False,
+         tab_width=4):
+    entry = _make_entry(cycles, case)
+    return format_line(line, entry, total, total_cond, subt, update,
+                       column, debug, use_tabs, tab_width)
+
+
+def test_adds_comment():
+    out = _run("  OPCODE", "5")
+    assert out == "  OPCODE      ; [5]\n"
+
+
+def test_subtotal_unconditional_opcode():
+    out = _run("  OPCODE", "5", total=42, total_cond=0, subt=True)
+    assert out == "  OPCODE      ; [5 .. 42]\n"
+
+
+def test_subtotal_conditional_opcode():
+    out = _run("  OPCODE", "7/5", total=42, total_cond=40, subt=True)
+    assert out == "  OPCODE      ; [7/5 .. 40/42]\n"
+
+
+def test_adds_case_if_debug_is_True():
+    out = _run("  OPCODE", "5", debug=True, case="foo")
+    assert out == "  OPCODE      ; [5] case{foo}\n"
+
+
+def test_adds_cycles_in_previous_comment_if_update_is_False():
+    out = _run("  OPCODE   ; [3] comment ", "5", update=False)
+    assert out == "  OPCODE   ; [5] [3] comment\n"
+
+
+def test_updates_cycles_in_previous_comment_if_update_is_True():
+    out = _run("  OPCODE   ; [3] comment ", "5", update=True)
+    assert out == "  OPCODE   ; [5] comment\n"
+
+
+def test_line_longer_than_comment_column():
+    out = _run("  A VERY LONG LINE", "5", update=True)
+    assert out == "  A VERY LONG LINE ; [5]\n"
diff --git a/tests/test_update_counters.py b/tests/test_update_counters.py
new file mode 100644
index 0000000..517e165
--- /dev/null
+++ b/tests/test_update_counters.py
@@ -0,0 +1,22 @@
+# -*- coding: utf-8 -*-
+
+from z80count.z80count import update_counters
+
+
+def _make_entry(states, states_met=0):
+    return {
+        "_t_states_met": states_met,
+        "_t_states_or_not_met": states,
+    }
+
+
+def test_unconditional_instruction():
+    total, total_cond = update_counters(_make_entry(3), 8)
+    assert total == 11
+    assert total_cond == 0
+
+
+def test_conditional_instruction():
+    total, total_cond = update_counters(_make_entry(7, 5), 35)
+    assert total == 42
+    assert total_cond == 40
diff --git a/z80count/z80count.py b/z80count/z80count.py
index 697450d..c59b0e1 100644
--- a/z80count/z80count.py
+++ b/z80count/z80count.py
@@ -35,49 +35,117 @@ OUR_COMMENT = re.compile(r"(\[[0-9.\s/]+\])")
 def z80count(line,
              parser,
              total,
-             total_cond,
              subt,
              no_update,
-             tabstop=2,
+             column=50,
+             use_tabs=False,
+             tab_width=8,
              debug=False,
              ):
     out = line.rstrip() + "\n"
     entry = parser.lookup(line)
     if entry:
-        cycles = entry["cycles"]
-        if "/" in cycles:
-            c = cycles.split("/")
-            total_cond = total + int(c[0])
-            total = total + int(c[1])
+        total, total_cond = update_counters(entry, total)
+        out = format_line(
+            line, entry, total, total_cond, subt, update=not no_update,
+            column=column, debug=debug, use_tabs=use_tabs,
+            tab_width=tab_width,
+        )
+    return (out, total)
+
+
+def update_counters(entry, total):
+    if entry["_t_states_met"]:
+        total_cond = total + entry["_t_states_met"]
+    else:
+        total_cond = 0
+    total = total + entry["_t_states_or_not_met"]
+
+    return (total, total_cond)
+
+
+def format_line(line, entry, total, total_cond, subt, update, column,
+                debug, use_tabs, tab_width):
+    cycles = entry["cycles"]
+    line = line.rstrip().rsplit(";", 1)
+    comment = "; [%s" % cycles
+    if subt:
+        if total_cond:
+            comment += " .. %d/%d]" % (total_cond, total)
         else:
-            total = total + int(cycles)
-            total_cond = 0
-
-        line = line.rstrip().rsplit(";", 1)
-        comment = "; [%s" % cycles
-        if subt:
-            if total_cond:
-                comment += " .. %d/%d]" % (total_cond, total)
-            else:
-                comment += " .. %d]" % total
+            comment += " .. %d]" % total
+    else:
+        comment += "]"
+    if debug:
+        comment += " case{%s}" % entry["case"]
+
+    if len(line) == 1:
+        comment = comment_alignment(line[0], column, use_tabs, tab_width) + comment
+    out = line[0] + comment
+    if len(line) > 1:
+        if update:
+            m = OUR_COMMENT.search(line[1])
+            if m:
+                line[1] = line[1].replace(m.group(0), "")
+        out += " "
+        out += line[1].lstrip()
+    out += "\n"
+
+    return out
+
+
+def comment_alignment(line, column, use_tabs=False, tab_width=8):
+    """Calculate the spacing required for comment alignment.
+
+    :param str line: code line
+    :param int column: column in which we want the comment to start
+    :param bool use_tabs: use tabs instead of spaces
+    :param int tab_width: tab width
+
+    :returns: the spacing
+    :rtype: str
+
+    """
+
+    expected_length = column - 1
+    length = line_length(line, tab_width)
+    if length >= expected_length:
+        return " "  # add an space before the colon
+
+    if use_tabs:
+        tab_stop = (expected_length // tab_width) * tab_width + 1
+        if tab_stop > length:
+            extra_tabs = (tab_stop - length) // tab_width
+            if length % tab_width > 1:
+                extra_tabs += 1  # complete partial tab
+            extra_spaces = expected_length - tab_stop
         else:
-            comment += "]"
-        if debug:
-            comment += " case{%s}" % entry["case"]
+            extra_tabs = 0
+            extra_spaces = expected_length - length
+    else:
+        extra_tabs = 0
+        extra_spaces = expected_length - length
+
+    return "\t" * extra_tabs + " " * extra_spaces
 
-        if len(line) == 1:
-            comment = "\t" * tabstop + comment
-        out = line[0] + comment
-        if len(line) > 1:
-            if not no_update:
-                m = OUR_COMMENT.search(line[1])
-                if m:
-                    line[1] = line[1].replace(m.group(0), "")
-            out += " "
-            out += line[1].lstrip()
-        out += "\n"
 
-    return (out, total, total_cond)
+def line_length(line, tab_width):
+    """Calculate the length of a line taking TABs into account.
+
+    :param str line: line of code
+    :param int tab_width: tab width
+
+    :returns: The length of the line
+    :rtype: int
+
+    """
+    length = 0
+    for i in line:
+        if i == "\t":
+            length = ((length + tab_width) // tab_width) * tab_width
+        else:
+            length += 1
+    return length
 
 
 def parse_command_line():
@@ -93,8 +161,15 @@ def parse_command_line():
                         help="Include subtotal")
     parser.add_argument('-n', dest='no_update', action='store_true',
                         help="Do not update existing count if available")
-    parser.add_argument('-t', dest='tabstop', type=int,
-                        help="Number of tabs for new comments", default=2)
+    parser.add_argument('-t', dest='tab_width', type=int,
+                        help="Number of spaces for each tab", default=8)
+    parser.add_argument('--use-spaces', dest='use_spaces', action='store_true',
+                        help="Use spaces to align newly added comments", default=True)
+    parser.add_argument('--use-tabs', dest='use_spaces', action='store_false',
+                        help="Use tabs to align newly added comments")
+    parser.add_argument('-c', '--column', dest='column', type=int,
+                        help="Column to align newly added comments", default=50)
+
     parser.add_argument(
         "infile", nargs="?", type=argparse.FileType('r'), default=sys.stdin,
         help="Input file")
@@ -122,9 +197,8 @@ class Parser(object):
         if mnemo is None or mnemo not in self._table:
             return None
         for entry in self._table[mnemo]:
-            if "cregex" not in entry:
-                entry["cregex"] = re.compile(
-                    r"^\s*" + entry["regex"] + r"\s*(;.*)?$", re.I)
+            if "_initialized" not in entry:
+                self._init_entry(entry)
             if entry["cregex"].search(line):
                 return entry
         return None
@@ -153,15 +227,32 @@ class Parser(object):
             return match.group("operator").upper()
         return None
 
+    @staticmethod
+    def _init_entry(entry):
+        entry["cregex"] = re.compile(r"^\s*" + entry["regex"] + r"\s*(;.*)?$", re.I)
+        cycles = entry["cycles"]
+        if "/" in cycles:
+            c = cycles.split("/")
+            t_states_or_not_met = int(c[1])
+            t_states_met = int(c[0])
+        else:
+            t_states_or_not_met = int(cycles)
+            t_states_met = 0
+        entry["_t_states_or_not_met"] = t_states_or_not_met
+        entry["_t_states_met"] = t_states_met
+        entry["_initialized"] = True
+
 
 def main():
     args = parse_command_line()
     in_f = args.infile
     out_f = args.outfile
     parser = Parser()
-    total = total_cond = 0
+    total = 0
     for line in in_f:
-        output, total, total_cond = z80count(
-            line, parser, total, total_cond,
-            args.subt, args.no_update, args.tabstop, args.debug)
+        output, total = z80count(
+            line, parser, total, args.subt, args.no_update,
+            args.column, not args.use_spaces, args.tab_width,
+            args.debug,
+        )
         out_f.write(output)