#!/usr/bin/env python3

'''
Tokenize disassembly from objdump.
'''

from   signal import signal, SIGPIPE, SIG_DFL
import sys

# Initial tokens.
INITIAL_TOKENS = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9', '0x', ' ', ',', 'a', 'b', 'c', 'd', 'e', 'f']

# How many records to process at a time.
BATCH_SIZE = 1000000

def handle_token(tok):
    '''
    Split a token into subtokens.
    :param tok: A string containing one or more tokens.
    :returns: The set of unique subtokens of tok.
    '''
    if tok.isalpha():
        return {tok}
    tokens = set()
    for c in tok:
        if c.isprintable() and not c.isalnum():
            tokens.add(c)
            subtoks = tok.split(c)
            for subtok in subtoks:
                tokens = tokens.union(handle_token(subtok))
    return tokens

if __name__ == '__main__':
    # Ignore closed pipe error.
    signal(SIGPIPE, SIG_DFL)
    tokens = set(INITIAL_TOKENS)
    line_num = 1
    print('{}: Processing input'.format(sys.argv[0]), file=sys.stderr)
    for line in sys.stdin:
        # Break into opcode and disassembly.
        elems = line.split('|')
        if len(elems) != 2:
            print('{}: Warning: {}: Expected 2 columns, got {}'.format(
                sys.argv[0],
                line_num,
                len(elems),
                file=sys.stderr
            ))
            continue
        disasm = elems[1]
        # Break disassembly into instruction and arguments.
        elems = disasm.replace(',', ' ').split()
        for tok in elems:
            tokens = tokens.union(handle_token(tok))
        line_num += 1
        # Show batch completed.
        if line_num % BATCH_SIZE == 0:
            print(
                '{}: Processed {} records and found {} unique tokens'.format(
                    sys.argv[0],
                    line_num,
                    len(tokens)
                ),
                file=sys.stderr
            )
    # Print sorted tokens.
    print('{}: Printing {} tokens'.format(sys.argv[0], len(tokens)), file=sys.stderr)
    for tok in sorted(tokens):
        print(tok)
