text2asm.py
mdsteele_annalog/build/text2asm.py
#=============================================================================#
# Copyright 2022 Matthew D. Steele <mdsteele@alum.mit.edu>                    #
#                                                                             #
# This file is part of Annalog.                                               #
#                                                                             #
# Annalog is free software: you can redistribute it and/or modify it under    #
# the terms of the GNU General Public License as published by the Free        #
# Software Foundation, either version 3 of the License, or (at your option)   #
# any later version.                                                          #
#                                                                             #
# Annalog is distributed in the hope that it will be useful, but WITHOUT ANY  #
# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS   #
# FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more      #
# details.                                                                    #
#                                                                             #
# You should have received a copy of the GNU General Public License along     #
# with Annalog.  If not, see <http://www.gnu.org/licenses/>.                  #
#=============================================================================#

import os
import sys

#=============================================================================#

MAX_PAIRS = 0xfd - 0x80

HEADER = """\
;;; This file was generated by text2asm.

.INCLUDE "../../src/charmap.inc"
.INCLUDE "../../src/dialog.inc"

;;;=========================================================================;;;

"""

FOOTER = """\

;;;=========================================================================;;;
"""

#=============================================================================#

def parse_text(data):
    text = []
    current_chars = []
    original_size = 0
    def finish_string():
        if not current_chars: return
        string = ''.join(current_chars)
        current_chars.clear()
        text.append(('b', string))
    while data:
        if data.startswith('{'):
            i = data.find('}')
            finish_string()
            constant = data[1:i]
            data = data[i + 1:]
            text.append(('c', constant))
            original_size += 1
        elif data.startswith('['):
            i = data.find(']')
            finish_string()
            pair = data[1:i]
            data = data[i + 1:]
            assert len(pair) == 4
            text.append(('p', (int(pair[:2], 16), int(pair[2:], 16))))
            original_size += 2
        else:
            char = data[0]
            data = data[1:]
            current_chars.append(char)
            original_size += 1
    finish_string()
    return (text, original_size)

def parse_input_file(filepath):
    texts = {}
    current_text_name = None
    current_text_data = ''
    original_size = 0
    for line in open(filepath):
        line = line.rstrip('\n')
        if current_text_name is not None:
            current_text_data += line
            if line.endswith('#') or line.endswith('%'):
                (text, size) = parse_text(current_text_data)
                texts[current_text_name] = text
                original_size += size
                current_text_name = None
                current_text_data = ''
            else:
                current_text_data += '$'
        else:
            if not line: continue
            elif line.startswith('#'): continue
            elif line.startswith('@'):
                current_text_name = line[1:]
            else:
                raise ValueError('bad line: ' + repr(line))
    assert current_text_name is None
    return (texts, original_size)

def compute_pairs(texts):
    forced_pairs = set()
    pair_counts = {}
    for text in texts.values():
        for kind, value in text:
            if kind == 'p':
                forced_pairs.add(value)
            elif kind == 'b':
                for i in range(0, len(value) - 1):
                    pair = value[i:i + 2]
                    if pair not in pair_counts:
                        pair_counts[pair] = 0
                    pair_counts[pair] += 1
    sorted_counts = sorted(pair_counts.items(), key=lambda item: -item[1])
    best_pairs = [pair for pair, count in sorted_counts if count > 2]
    return sorted(forced_pairs) + best_pairs[:MAX_PAIRS - len(forced_pairs)]

def compress_text(text, dictionary):
    compressed_size = 0
    result_lines = []
    current_line = []
    current_chars = []
    def finish_string():
        if not current_chars: return
        string = ''.join(current_chars)
        current_chars.clear()
        current_line.append(f'"{string}"')
    def finish_line():
        finish_string()
        if not current_line: return
        line = ', '.join(current_line)
        current_line.clear()
        result_lines.append(f'    .byte {line}\n')
    for kind, value in text:
        if kind == 'b':
            while len(value) >= 2:
                pair = value[:2]
                i = dictionary.get(pair)
                if i is not None:
                    value = value[2:]
                    finish_string()
                    current_line.append(f'${i + 0x80:02x}')
                    compressed_size += 1
                    if '$' in pair: finish_line()
                else:
                    char = value[0]
                    value = value[1:]
                    current_chars.append(char)
                    compressed_size += 1
                    if char == '$': finish_line()
            if value:
                current_chars.append(value)
                compressed_size += len(value)
                if value.endswith('$'): finish_line()
        elif kind == 'c':
            finish_string()
            current_line.append(value)
            compressed_size += 1
        elif kind == 'p':
            i = dictionary[value]
            finish_string()
            current_line.append(f'${i + 0x80:02x}')
            compressed_size += 1
        else: assert False
    finish_line()
    return (''.join(result_lines), compressed_size)

def write_output_file(bank, pairs, texts, original_data_size):
    compressed_data_size = 0
    sys.stdout.write(HEADER)
    sys.stdout.write(f'.SEGMENT "PRGA_{bank}"\n\n')
    sys.stdout.write(f'.EXPORT DataA_{bank}_Strings_u8_arr2_arr\n')
    sys.stdout.write(f'.PROC DataA_{bank}_Strings_u8_arr2_arr\n')
    for pair in pairs:
        if isinstance(pair, str):
            sys.stdout.write(f'    .byte "{pair}"\n')
        else:
            sys.stdout.write(f'    .byte ${pair[0]:02x}, ${pair[1]:02x}\n')
        compressed_data_size += 2
    sys.stdout.write(f'.ENDPROC\n')
    dictionary = {pair: i for i, pair in enumerate(pairs)}
    for name, text in sorted(texts.items()):
        sys.stdout.write(f'\n.EXPORT DataA_{bank}_{name}_u8_arr\n')
        sys.stdout.write(f'.PROC DataA_{bank}_{name}_u8_arr\n')
        (compressed_asm, compressed_size) = compress_text(text, dictionary)
        sys.stdout.write(compressed_asm)
        sys.stdout.write(f'.ENDPROC\n')
        compressed_data_size += compressed_size
    sys.stdout.write(f'\n;;;   Original size = {original_data_size:4x}\n')
    sys.stdout.write(f';;; Compressed size = {compressed_data_size:4x}\n')
    saved = original_data_size - compressed_data_size
    sys.stdout.write(f';;;     Bytes saved = {saved:4x}\n')
    percent = int(round(100 * (1 - compressed_data_size / original_data_size)))
    sys.stdout.write(f';;;   Percent saved = {percent:3d}%\n')
    sys.stdout.write(FOOTER)

def run(filepath):
    filename = os.path.split(filepath)[1]
    bank = os.path.splitext(filename)[0].capitalize()
    (texts, original_size) = parse_input_file(filepath)
    pairs = compute_pairs(texts)
    write_output_file(bank, pairs, texts, original_size)

#=============================================================================#

if __name__ == '__main__':
    run(sys.argv[1])

#=============================================================================#
