Fix coverage utility script

3 files changed, 41 insertions, 37 deletions

diff --git a/Makefile b/Makefile
index b546677..f3a0f4b 100644
--- a/Makefile
+++ b/Makefile
@@ -57,8 +57,8 @@ tidy:
 
 # Visualize disassembly progress.
 .PHONY: coverage
-coverage: $(ROM:.gb=.map) tools/disasm_coverage.py
-	$(PYTHON) tools/disasm_coverage.py -m $< -b 0x40
+coverage: $(ROM:.gb=.map) utils/disasm_coverage.py
+	$(PYTHON) utils/disasm_coverage.py -m $< -b 0x40
 
 
 %.map: %.gb
diff --git a/tools/disasm_coverage.py b/utils/disasm_coverage.py
index fcb9e86..6a68730 100644
--- a/tools/disasm_coverage.py
+++ b/utils/disasm_coverage.py
@@ -45,7 +45,8 @@ if __name__ == '__main__':
     if (romname is not None):
         with open(romname, 'rb') as f:
             for rb in range(0, args.num_banks):
-                data = r.bank_data['ROM Bank'].get(rb, default_bank_data)
+                bank_data = r.bank_data['ROM0 bank' if rb == 0 else 'ROMX bank']
+                data = bank_data.get(rb, default_bank_data)
                 bank = f.read(bank_size)
                 if (bank[bank_size - 1] in filler):
                     fill = bank[bank_size - 1]
@@ -59,7 +60,8 @@ if __name__ == '__main__':
     
     hit_data = [[0] * width for _ in range(height)]
     for bank in range(args.num_banks):
-        data = r.bank_data['ROM Bank'].get(bank, default_bank_data)
+        bank_data = r.bank_data['ROM0 bank' if bank == 0 else 'ROMX bank']
+        data = bank_data.get(bank, default_bank_data)
         for s in data['sections']:
             beg = (s['beg'] & bank_mask) + bank * bank_size
             end = ((s['end'] -1) & bank_mask) + bank * bank_size # end is exclusive
diff --git a/tools/mapreader.py b/utils/mapreader.py
index 95a4723..1164cc4 100644
--- a/tools/mapreader.py
+++ b/utils/mapreader.py
@@ -1,6 +1,7 @@
-#!/usr/bin/env python3
 # -*- coding: utf-8 -*-
 
+# A library for parsing the pokegold-spaceworld.map file output by rgbds.
+
 import re
 
 class MapReader:
@@ -37,30 +38,31 @@ class MapReader:
     bank_types = {
         'HRAM'     : { 'size': 0x80,   'banked': False, },
         'OAM'      : { 'size': 0xA0,   'banked': False, },
-        'ROM Bank' : { 'size': 0x4000, 'banked': True,  },
-        'SRAM Bank': { 'size': 0x2000, 'banked': True,  },
-        'VRAM Bank': { 'size': 0x1000, 'banked': True,  },
-        'WRAM Bank': { 'size': 0x2000, 'banked': True,  },
+        'ROM0 bank': { 'size': 0x4000, 'banked': True,  },
+        'ROMX bank': { 'size': 0x4000, 'banked': True,  },
+        'SRAM bank': { 'size': 0x2000, 'banked': True,  },
+        'VRAM bank': { 'size': 0x1000, 'banked': True,  },
+        'WRAM bank': { 'size': 0x2000, 'banked': True,  },
     }
-    
+
     # FSM states
     INIT, BANK, SECTION = range(3)
-    
+
     # $506D-$519A ($012E bytes) ["Type Matchups"]
     section_header_regex = re.compile('\$([0-9A-Fa-f]{4})-\$([0-9A-Fa-f]{4}) \(.*\) \["(.*)"\]')
     # $506D = TypeMatchups
     section_data_regex = re.compile('\$([0-9A-Fa-f]{4}) = (.*)')
     # $3ED2 bytes
-    slack_regex = re.compile('\$([0-9A-Fa-f]{4}) bytes')
+    slack_regex = re.compile('\$([0-9A-Fa-f]{4}) bytes?')
 
     def __init__(self, *args, **kwargs):
         self.__dict__.update(kwargs)
-    
+
     def _parse_init(self, line):
-    
+
         line = line.split(':', 1)[0]
         parts = line.split(' #', 1)
-        
+
         if (parts[0] in self.bank_types):
             self._cur_bank_name = parts[0]
             self._cur_bank_type = self.bank_types[self._cur_bank_name]
@@ -76,17 +78,17 @@ class MapReader:
                 self._cur_data = self.bank_data[self._cur_bank_name][self._cur_bank]
             else:
                 self._cur_data = self.bank_data[self._cur_bank_name]
-            
+
             if ({} == self._cur_data):
                 self._cur_data['sections'] = []
                 self._cur_data['used'] = 0
                 self._cur_data['slack'] = self._cur_bank_type['size']
             return True
-        
+
         return False
-    
+
     def _parse_section_header(self, header):
-    
+
         section_data = self.section_header_regex.match(header)
         if section_data is not None:
             beg = int(section_data.group(1), 16)
@@ -96,48 +98,48 @@ class MapReader:
             self._cur_data['sections'].append(self._cur_section)
             return True
         return False
-        
+
     def _parse_slack(self, data):
-    
+
         slack_data = self.slack_regex.match(data)
         slack_bytes = int(slack_data.group(1), 16)
         self._cur_data['slack'] = slack_bytes
-        
+
         used_bytes = 0
-        
+
         for s in self._cur_data['sections']:
             used_bytes += s['end'] - s['beg'] + 1
-    
+
         self._cur_data['used'] = used_bytes
-    
+
     def read_map_data(self, map):
-        
+
         if type(map) is str:
             map = map.split('\n')
-        
+
         self._state = MapReader.INIT
         self._cur_bank_name = ''
         self._cur_bank_type = {}
         self._cur_bank = 0
         self._cur_data = {}
-        
+
         for line in map:
-            
+
             line = line.rstrip()
             if (MapReader.INIT == self._state):
-            
+
                 if (self._parse_init(line)):
                     self._state = MapReader.BANK
-            
+
             elif (MapReader.BANK == self._state or MapReader.SECTION == self._state):
-            
+
                 if ('' == line):
                     self._state = MapReader.INIT
                 else:
-                    
+
                     line = line.lstrip()
                     parts = line.split(': ', 1)
-                    
+
                     if (MapReader.SECTION == self._state):
                         section_data = self.section_data_regex.match(parts[0])
                         if section_data is not None:
@@ -145,7 +147,7 @@ class MapReader:
                             name = section_data.group(2)
                             self._cur_section['symbols'].append({'name': name, 'address': address})
                             continue
-                    
+
                     if ('SECTION' == parts[0]):
                         if (self._parse_section_header(parts[1])):
                             self._state = MapReader.SECTION
@@ -155,10 +157,10 @@ class MapReader:
                     elif ('EMPTY' == parts[0]):
                         self._cur_data = {'sections': [], 'used': 0, 'slack': self._cur_bank_type['size']}
                         self._state = MapReader.INIT
-            
+
             else:
                 pass
-       
+
         for k, v in self.bank_data.items():
             if (self.bank_types[k]['banked']):
                 for _, vv in v.items():