#!/usr/bin/python3
"""Search for the simplest way to compute all binary Boolean functions.

I wanted to see what the impact of different possible instruction sets
would be on the difficulty of expressing things.  How much of a pain
would it be to have just NAND?  How about just abjunction?  How about
NAND and also XOR?  (Turns out NOR requires three gates in that case,
but everything else can be done in two.)

This program lets you ask those questions; it shows you everything you
can compute from two inputs while minimizing some cost function.

The basic data structure is a list of (operation, [operands], results)
tuples.  We start with a minimal list [('x', [], 3), ('y', [], 5)] and
keep trying to add things to the list.  Once we attempt to apply all
the available ops to all the available sets of operands without
achieving a cost reduction, we are done.  If we find a cheaper way to
compute an existing result, we replace it.

The simplest way to compute the cost requires a graph traversal, since
the operands may have shared structure, which means we can’t simply
add the costs of computing them.

"""
import sys


ALL_ONE = 0xf

# XXX use a decorator

def abj(x, y):
    "Abjunction.  andnot.  bic."
    return x & ~y
abj.format = "{} &^ {}".format  # from Golang
abj.arity = 2

def not_(x):
    return ~x & ALL_ONE
not_.format = "~{}".format
not_.arity = 1

def and_(x, y):
    return x & y
and_.format = "{} & {}".format
and_.arity = 2

def or_(x, y):
    return x | y
or_.format = '{} | {}'.format
or_.arity = 2

def nand(x, y):
    return ~(x & y) & ALL_ONE
nand.format = "{} &̄ {}".format
nand.arity = 2

def nand3(x, y, z):
    return ~(x & y & z) & ALL_ONE
nand3.format = "nand({}, {}, {})".format
nand3.arity = 3

def nand4(w, x, y, z):
    return ~(w & x & y & z) & ALL_ONE
nand4.format = "nand({}, {}, {}, {})".format
nand4.arity = 4

def xor(x, y):
    return x ^ y
xor.format = "{} ^ {}".format
xor.arity = 2

def ifthen(x, y, z):
    return x & y | ~x & z
ifthen.format = '{} ? {} : {}'.format
ifthen.arity = 3

def notifthen(x, y, z):
    return ALL_ONE & ~(x & y | ~x & z)
notifthen.format = '~({} ? {} : {})'.format
notifthen.arity = 3

def aoi(a, b, c, d):
    "And-or-invert."
    return ALL_ONE & ~(a & b | c & d)
aoi.format = '~({} & {} | {} & {})'.format
aoi.arity = 4

def search(ops, debug=print, include_rails=False, costs=lambda op: 1):
    """Search for ways to produce logic functions of two inputs using ops.

    ops is a list of callables with the `arity` property set on them.
    debug is a function which produces debugging output.
    include_rails specifies whether all-ones and all-zeroes are available.
    costs is the cost function we’re trying to minimize the total cost of.

    This uses a pretty dumb algorithm, but for up to 16 entries it’s okay.
    """
    found = [('x', [], 3),
             ('y', [], 5)]
    if include_rails:          # y’know, like the power rails in a logic circuit
        found.extend([('0', [], 0),
                      ('-1', [], ALL_ONE)])

    improved = True
    while improved:
        improved = False
        for op in ops:
            for tuple_indices in index_sets(op.arity, range(len(found))):
                tuples = [found[i] for i in tuple_indices]
                results = op(*(result for _, _, result in tuples))
                cost = compute_cost(found, tuple_indices, op, costs)
                result_tuple = (op, tuple_indices, results)

                for i, old_tuple in enumerate(found):
                    (old_op, old_indices, result) = old_tuple
                    if result == results:
                        old_cost = compute_cost(found, old_indices, old_op, costs)
                        if old_cost > cost:
                            debug("replacing", old_tuple, "with", result_tuple)
                            found[i] = result_tuple
                            improved = True

                        # “A break statement executed in the first
                        # suite terminates the loop without executing
                        # the else clause’s suite.” — PLR 3.5 §8.3
                        break
                else:
                    debug("adding", result_tuple)
                    found.append(result_tuple)
                    improved = True

    return found

# XXX itertools.product
def index_sets(n, m):
    "Return all tuples of n items each selected from m."
    if n == 0:
        yield ()
        return

    for rest in index_sets(n-1, m):
        for i in m:
            yield rest + (i,)

def compute_cost(universe, start_indices, start_op, costs):
    """Compute the cost of applying op to the quantities at start_indices.

    This includes the cost of the quantities at start_indices, but it
    may be less than the sum of their costs, because some of their
    costs may be shared, for example if you compute nand(foo, foo).

    """
    cost = 0
    seen = set()
    pending_ops = [(start_op, start_indices)]
    
    while pending_ops:
        op, indices = pending_ops.pop()
        cost += costs(op)
        for index in indices:
            if index not in seen:
                seen.add(index)
                next_op, next_indices, _ = universe[index]
                pending_ops.append((next_op, next_indices))

    return cost

def cost(universe, i, costs):
    "Compute the cost of computing item i; see compute_cost for details."
    op, indices, _ = universe[i]
    return compute_cost(universe, indices, op, costs)


def infix(universe, i):
    """Render item i of universe in infix notation.

    Uses a where-clause if necessary.

    This function can probably be rewritten more simply.
    """
    refcounts = {}
    tovisit = [i]
    while tovisit:
        n = tovisit.pop()
        if n not in refcounts:
            refcounts[n] = 0
        refcounts[n] += 1
        _, indices, _ = universe[n]
        tovisit.extend(indices)

    named = [var for var in refcounts if refcounts[var] > 1]
    named_ord = sorted(named, reverse=True)
    names = {k: chr(ord('a') + j) for j, k in enumerate(named_ord)}

    # Cheat for x, y, 0, and 1:
    for j, (op, indices, _) in enumerate(universe):
        if not indices:
            names[j] = op

    top = infix2(universe, i, names)
    lets = [names[j] + " = " + infix2(universe, j, names)
            for j in named_ord
            if universe[j][1]]
    if lets:
        return top + " where " + " and ".join(lets)
    else:
        return top

def infix2(universe, i, names):
    "Internal recursive subfunction of infix."
    op, indices, _ = universe[i]
    if not indices:
        return op
    subexpressions = [names[j] if j in names else
                      '(%s)' % infix2(universe, j, names)
                      for j in indices]
    return op.format(*subexpressions)


def main(opnames):
    op_dict = {
        '&^': abj,
        '\\': abj,
        '~': not_,
        '&': and_,
        '|': or_,
        '&̄': nand,
        '^': xor,
        '⊕': xor,
        '?:': ifthen,
    }
    for op in [abj, not_, and_, or_, nand, nand3, nand4,
               xor, ifthen, notifthen, aoi]:
        op_dict[op.__name__] = op
        if op.__name__.endswith('_'): # XXX strip
            op_dict[op.__name__[:-1]] = op
        
    ops = ([op_dict[name] for name in opnames] if opnames else [abj])
    costs = lambda op: 0 if op == 'x' or op == 'y' else 1

    print("searching with ops:", *[op.__name__ for op in ops])
    universe = search(ops, debug=lambda *args: None, costs=costs)

    if len(universe) < 16:
        print("adding power rails")
        universe = search(ops, debug=lambda *args: None, costs=costs,
                          include_rails=True)

    for i, (op, indices, result) in enumerate(universe):
        if indices:
            expr = op.format(*("r[%d]" % i for i in indices))
        else:
            expr = op

        print("r[{:2d}] = {} (truth table {:04b}, cost {}) = {}".format(
                  i, expr, result, cost(universe, i, costs),
            infix(universe, i)))

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