1 files changed, 4 insertions, 5 deletions
diff --git a/format.py b/format.py
index d43f835..419f6df 100755
--- a/format.py
+++ b/format.py
@@ -283,7 +283,6 @@ instr_map = {
 }
 
 def format(iter):
-    segment = None
     for line_no, raw_line in enumerate(iter, start=1):
         try:
             line = subv.parse(raw_line)
@@ -296,7 +295,7 @@ parsed as {}
             raise Exception(message) from e
 
         try:
-            if segment == 'code' and line['type'] == 'instr':
+            if line['type'] == 'instr':
                 op = line['instr'][0]
                 assert len(op) == 2, 'instruction without op label: {}'.format(op)
 
@@ -308,11 +307,11 @@ parsed as {}
                     raise ValueError("opcode {} doesn't match label {} (expected {})"
                                      .format(op, label, expected))
 
-                line['instr'] = formatter(line['instr'])
+                formatted = formatter(line['instr'])
+                formatted[0] = (*formatted[0], label)
+                line['instr'] = formatted
                 yield subv.format(line)
             else:
-                if line['type'] == 'segment':
-                    segment = line['segment'][0]
                 yield line['raw']
         except AssertionError as e:
             message = '''