RadixMaker

#!/usr/bin/python2

import sys
import csv

def make_length_group(data):
    ldict = {}
    for s, v in data:
        ldict.setdefault(len(s), []).append((s, v))
    return ldict

def make_prefix_tree(data, eos=True, compact=False):
    def build_tree(data, eos):
        cdict = {}
        for st, val in data:
            if eos:
                if len(st) > 0:
                    cdict.setdefault(st[0], []).append((st[1:], val))
                else:
                    cdict[0] = val
            else:
                if len(st) > 1:
                    cdict.setdefault(st[0], []).append((st[1:], val))
                else:
                    cdict[st[0]] = val
        rdict = {}
        for key, value in cdict.iteritems():
            if isinstance(value, list):
                rdict[key] = build_tree(value, eos)
            else:
                rdict[key] = value
        return rdict

    def make_compact(node):
        if len(node) == 1 and isinstance(node.values()[0], dict) and len(node.values()[0]) == 1:
            cnode = node.values()[0]
            if cnode.keys()[0] == 0:
                newval = cnode[cnode.keys()[0]]
                node[node.keys()[0]] = newval
            else:
                newkey = node.keys()[0] + cnode.keys()[0]
                newval = cnode[cnode.keys()[0]]
                del node[node.keys()[0]]
                node[newkey] = newval
                make_compact(node)
        else:
            for key, value in node.iteritems():
                if isinstance(value, dict):
                    make_compact(value)

    node = build_tree(data, eos)
    if compact:
        make_compact(node)
    return node

def make_strcmp(key, index = 0, usemem = False):
    if isinstance(key, str):
        if usemem:
            return "!memcmp(s+%d, \"%s\", %d)" % (index, key, len(key))
        else:
            return "!strcmp(s+%d, \"%s\")" % (index, key)
    else:
        return "s[%d] == %d" % (index, key)

def make_strncmp(key, index = 0, usemem = False):
    if isinstance(key, str):
        if len(key) == 1:
            return "s[%d] == '%s'" % (index, key)
        else:
            if usemem:
                return "!memcmp(s+%d, \"%s\", %d)" % (index, key, len(key))
            else:
                return "!strncmp(s+%d, \"%s\", %d)" % (index, key, len(key))
    else:
        return "s[%d] == %d" % (index, key)

def make_prefix_tree_to_switch(data, depth = 0):
    ldict = make_length_group(data)
    t = make_prefix_tree(data, eos = True, compact = True)
    make_prefix_tree_to_switch_sub(t, depth, usemem = False)

def make_prefix_tree_to_switch_len(data, depth = 0):
    ldict = make_length_group(data)
    istr = '\t'*depth
    print "%sswitch (len) {" % (istr)
    for l, ldata in sorted(ldict.iteritems(), key = lambda x: x[0]):
        print "%scase %d:" % (istr, l),
        t = make_prefix_tree(ldata, eos = False, compact = True)
        make_prefix_tree_to_switch_sub(t, depth + 1, usemem = True)
    print "%sdefault: return -1;" % (istr)
    print "%s}" % (istr)

def make_prefix_tree_to_switch_sub(node, depth = 0, usemem = False, index = 0):
    istr = '\t'*depth
    if len(node) > 1:
        print
        print "%sswitch (s[%d]) {" % (istr, index)
        for key, value in sorted(node.iteritems(), key = lambda x: x[0]):
            if isinstance(key, str):
                print "%scase '%s':" % (istr, key),
            else:
                print "%scase %d:" % (istr, key),
            if isinstance(value, dict):
                make_prefix_tree_to_switch_sub(value, depth+1, usemem, index+1)
            else:
                print "return %d;" % (value);
        print "%sdefault: return -1;" % (istr)
        print "%s}" % (istr)
    else:
        key = node.keys()[0]
        val = node.values()[0]
        if isinstance(val, dict):
            print "if (%s)" % (make_strncmp(key, index, usemem)),
            make_prefix_tree_to_switch_sub(val, depth, usemem, index+len(key))
            print "%selse return -1;" % (istr);
        else:
            print "if (%s)" % (make_strcmp(key, index, usemem)),
            print "return %d; else return -1;" % (val)

# >copy con Test.txt
# apple 1
# apps  2
# bana  3
# banana    4

def main():
    # read input
    argv = sys.argv[1:]
    input = argv[0] if len(argv) > 0 else "Test.txt"
    data = []
    for s, v in csv.reader(open(input), csv.excel_tab):
        data.append((s, int(v)))
    # work
    print "int parse(const char* s) {",
    make_prefix_tree_to_switch(data, 1)
    print "}"
    print
    print "int parse(const char* s, size_t len) {"
    make_prefix_tree_to_switch_len(data, 1)
    print "}"

if __name__ == "__main__":
    main()