4-variable-simplifier/spikes/assignment.py

#!/usr/bin/env python3

from heapq import heappush, heappop

X = 0b01010101
Y = 0b00110011
Z = 0b00001111
M = 0b11111111

B = 1 << 3
N = 1 << B

def calc_simps(A = None):
    # truthtable -> ('var', name)
    # truthtable -> ('not', inner truthtable)
    # truthtable -> (binary op, left truthtable, right truthtable)
    lookup = dict();

    # truthtable -> cost
    costs = dict();

    todo = set();

    def initial(exp, truthtable):
        lookup[truthtable] = exp;
        costs[truthtable] = 0;
        todo.add(truthtable);

    initial(('var', 'x'), X);
    initial(('var', 'y'), Y);
    initial(('var', 'z'), Z);
    initial(('lit', '0'), 0);
    initial(('lit', '1'), M);
    
    if (A is not None):
        initial(('var', 'a'), A);
    
    done = list();

    for iteration in range(1, N + 1):
        mytruthtable = min(todo, key = lambda x: (costs[x], x));
        
        mycost = costs[mytruthtable];
        
        todo.remove(mytruthtable);
        
        # consider NOT:
        not_cost = mycost + 1
        not_truthtable = ~mytruthtable & M;
        if (not_truthtable not in costs) or (not_cost < costs[not_truthtable]):
            todo.add(not_truthtable);
            
            lookup[not_truthtable] = ("not", mytruthtable);
            
            costs[not_truthtable] = not_cost;
        
        # consider OR:
        for othertruthtable in done:
            or_cost = 1 + mycost + costs[othertruthtable];
            
            or_truthtable = mytruthtable | othertruthtable;
            
            if (or_truthtable not in costs) or (or_cost < costs[or_truthtable]):
                todo.add(or_truthtable);
                
                lookup[or_truthtable] = ("or", mytruthtable, othertruthtable);
                
                costs[or_truthtable] = or_cost;
        
        # consider AND:
        for othertruthtable in done:
            and_cost = 1 + mycost + costs[othertruthtable];
            
            and_truthtable = mytruthtable & othertruthtable;
            
            if (and_truthtable not in costs) or (and_cost < costs[and_truthtable]):
                todo.add(and_truthtable);
                
                lookup[and_truthtable] = ("and", mytruthtable, othertruthtable);
                
                costs[and_truthtable] = and_cost;
        
        # consider XOR:
        if 0:
            for othertruthtable in done:
                xor_cost = 1 + mycost + costs[othertruthtable];
                
                xor_truthtable = mytruthtable ^ othertruthtable;
                
                if (xor_truthtable not in costs) or (xor_cost < costs[xor_truthtable]):
                    todo.add(xor_truthtable);
                    
                    lookup[xor_truthtable] = ("xor", mytruthtable, othertruthtable);
                    
                    costs[xor_truthtable] = xor_cost;
        
        done.append(mytruthtable);

    assert(sorted(done) == sorted(costs));
    
    return costs, lookup;

def etostr(lookup, expr):
    match expr:
        case ('lit', x):
            return x;
        case ('var', x):
            return x;
        case ('not', inner):
            return "(!" + tostr(lookup, inner) + ")";
        case ('or', left, right):
            return "(" + tostr(lookup, left) + " || " + tostr(lookup, right) + ")";
        case ('and', left, right):
            return "(" + tostr(lookup, left) + " && " + tostr(lookup, right) + ")";
        case ('xor', left, right):
            return "(" + tostr(lookup, left) + " != " + tostr(lookup, right) + ")";
        case _:
            assert(not "TODO");

def tostr(lookup, truthtable):
    return etostr(lookup, lookup[truthtable]);

zdcosts, zdlookup = calc_simps();

# now that we've figured out the zero-depth trees, we need to figure out which
# trees could be done faster with a variable. We need to go through every
# possible truthtable, considering them the new variable, building up a new
# table of what expressions are the cheapiest with that variable.
# Then, if that figured out for all of them, we could read down the columns
# figuring out the best answers for anything the user gives.

# truthtable -> cost
costs = zdcosts;

# truthtable -> variable truthtable
varlookup = {x: None for x in range(1 << B)};

all_costs = dict(); # variable truthtable -> truthtable -> cost
all_lookups = dict(); # variable truthtable -> truthtable -> expr

for A in range(0, 1 << B):
    print(f"what if the variable was 0b{A:08b}?");
    
    var_costs, var_lookup = calc_simps(A);
    
    for t in range(0, 1 << B):
        mycost = var_costs[t] + 1 + zdcosts[A];
        
        if mycost < costs[t]:
            varlookup[t] = A;
            
            costs[t] = mycost;
    
    all_costs[A] = var_costs;
    
    all_lookups[A] = var_lookup;
    
for t in range(0, 1 << B):
    line = f"0b{t:08b}: ";
    
    line += f"[{costs[t]}] ";
    
    if varlookup[t] is None:
        line += tostr(zdlookup, t);
    else:
        line += "("
        
        a = varlookup[t];
        
        line += "a = " + tostr(zdlookup, a);
        
        line += ", "
        
        line += tostr(all_lookups[a], t);
        
        line += ")"
        
        line += " [before: " + tostr(zdlookup, t) + "]";
    
    print(line);

print();
print(f"we can do anything in {max(costs.values())} operations");