boolean-tools/anf.py

#!/bin/python3
from __future__ import annotations

import re
from dataclasses import dataclass
from typing import FrozenSet, Iterable, List, Union

import sys

#####################
# ANF representation over GF(2)
#####################
#
# polynomial: set of monomials.
# monomial: frozenset of variable names.
#
# Examples:
#   0      -> set()
#   1      -> {frozenset()}
#   X3     -> {frozenset({"X3"})}
#   X3 XOR X4 -> {{"X3"}, {"X4"}}
#   X3 AND X4 -> {{"X3","X4"}}


Monomial = FrozenSet[str]
Poly = set[Monomial]


def zero() -> Poly:
    return set()


def one() -> Poly:
    return {frozenset()}


def var(name: str) -> Poly:
    return {frozenset([name])}


def poly_xor(a: Poly, b: Poly) -> Poly:
    # symmetric difference = XOR over GF(2)
    return a ^ b


def poly_and(a: Poly, b: Poly) -> Poly:
    out: Poly = set()
    for ma in a:
        for mb in b:
            m = frozenset(set(ma) | set(mb))
            if m in out:
                out.remove(m)
            else:
                out.add(m)
    return out


def poly_not(a: Poly) -> Poly:
    # NOT a = a XOR 1
    return poly_xor(a, one())


#####################
# Tokenizer
#####################

TOKEN_RE = re.compile(
    r"""
    \s*(
        \(|\)|
        !|
        AND|OR|
        0|1|
        X\d+|
        [A-Za-z_][A-Za-z0-9_]*
    )
    """,
    re.VERBOSE,
)

def tokenize(s: str) -> List[str]:
    tokens = TOKEN_RE.findall(s)
    stripped = re.sub(r"\s+", "", s)
    joined = "".join(tokens).replace("AND", "AND").replace("OR", "OR")
    if re.sub(r"\s+", "", joined) != stripped:
        raise ValueError(f"Invalid token near: {s!r}")
    return tokens


#####################
# Parser
#####################
# priotity order:
#   !
#   AND
#   OR
#
# Grammar:
#   expr   := or_expr
#   or_expr := and_expr ("OR" and_expr)*
#   and_expr := not_expr ("AND" not_expr)*
#   not_expr := "!" not_expr | atom
#   atom := VAR | 0 | 1 | "(" expr ")"


class Parser:
    def __init__(self, tokens: List[str]):
        self.tokens = tokens
        self.i = 0

    def peek(self) -> str | None:
        return self.tokens[self.i] if self.i < len(self.tokens) else None

    def eat(self, expected: str | None = None) -> str:
        tok = self.peek()
        if tok is None:
            raise ValueError("Unexpected end of input")
        if expected is not None and tok != expected:
            raise ValueError(f"Expected {expected!r}, got {tok!r}")
        self.i += 1
        return tok

    def parse(self) -> Poly:
        p = self.parse_or()
        if self.peek() is not None:
            raise ValueError(f"Unexpected token: {self.peek()!r}")
        return p

    def parse_or(self) -> Poly:
        left = self.parse_and()
        while self.peek() == "OR":
            self.eat("OR")
            right = self.parse_and()
            # A OR B = A XOR B XOR (A AND B)
            left = poly_xor(poly_xor(left, right), poly_and(left, right))
        return left

    def parse_and(self) -> Poly:
        left = self.parse_not()
        while self.peek() == "AND":
            self.eat("AND")
            right = self.parse_not()
            left = poly_and(left, right)
        return left

    def parse_not(self) -> Poly:
        if self.peek() == "!":
            self.eat("!")
            return poly_not(self.parse_not())
        return self.parse_atom()

    def parse_atom(self) -> Poly:
        tok = self.peek()
        if tok is None:
            raise ValueError("Unexpected end of input")

        if tok == "(":
            self.eat("(")
            inside = self.parse_or()
            self.eat(")")
            return inside

        if tok == "0":
            self.eat("0")
            return zero()

        if tok == "1":
            self.eat("1")
            return one()

        if re.fullmatch(r"X\d+|[A-Za-z_][A-Za-z0-9_]*", tok):
            self.eat()
            return var(tok)

        raise ValueError(f"Unexpected token: {tok!r}")

#####################
# CLI PRETTY PRINT
#####################

def monomial_key(m: Monomial):
    return (len(m), tuple(sorted(m)))


def format_monomial(m: Monomial) -> str:
    if len(m) == 0:
        return "1"
    if len(m) == 1:
        return next(iter(m))
    parts = sorted(m, key=lambda x: (len(x), x))
    return "(" + "⸱".join(parts) + ")"


def format_poly(p: Poly) -> str:
    if not p:
        return "0"

    monos = sorted(p, key=monomial_key)
    return "⊕".join(format_monomial(m) for m in monos)


def convert(expr: str) -> str:
    tokens = tokenize(expr)
    poly = Parser(tokens).parse()
    return format_poly(poly)

#####################
# CLI
#####################

if __name__ == "__main__":

    if len(sys.argv) > 1:
        expr = " ".join(sys.argv[1:])
    else:
        expr = input("Enter formula: ").strip()

    print(convert(expr))