Binary DRAT parser

parse_binary_drat.py

import mmap
import tqdm

def parse_binary_drat_mmap(path):
    with open(path, "rb") as f:
        with mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ) as mm:
            state = 0
            # state = 0 -- reading the mode: b'a' or b'd'
            # state = 1 -- reading the clause
            mode = None  # "a" or "d", for user
            lit = 0
            shift = 0
            clause = []

            for b in mm:
                if state == 0:
                    # Begin parsing a clause
                    if b == b"a":
                        mode = "a"
                    elif b == b"d":
                        mode = "d"
                    else:
                        raise ValueError(f"Bad clause header: {b}")
                    clause = []
                    lit = 0
                    shift = 0
                    state = 1

                elif state == 1:
                    lit |= (ord(b) & 127) << shift
                    shift += 7
                    if ord(b) <= 127:
                        # Finish parsing a literal
                        if lit % 2:
                            lit = -(lit >> 1)
                        else:
                            lit = lit >> 1

                        if lit == 0:
                            # Finish parsing a clause
                            yield (mode, clause)
                            mode = None
                            state = 0
                        else:
                            clause.append(lit)
                            # Reset parsing a literal
                            lit = 0
                            shift = 0

                else:
                    raise ValueError(f"Bad state: {state}")


def parse_binary_drat(path):
    def read_lit(file):
        lit = 0
        shift = 0
        while True:
            lc = file.read(1)
            if not lc:
                return
            lit |= (ord(lc) & 127) << shift
            shift += 7
            if ord(lc) <= 127:
                break

        if lit % 2:
            return -(lit >> 1)
        else:
            return lit >> 1

    with open(path, "rb") as f:
        while True:
            b = f.read(1)

            # Handle EOF:
            if not b:
                return

            # Determine whether the clause was "added" or "deleted":
            if b == b"a":
                mode = "a"
            elif b == b"d":
                mode = "d"
            else:
                raise ValueError(f"Bad clause header: {mode}")

            # Read the clause:
            clause = []
            while True:
                lit = read_lit(f)
                if lit == 0:
                    # End of the clause
                    break
                elif lit is None:
                    raise ValueError("Could not read literal")
                else:
                    # Another literal in the clause
                    clause.append(lit)

            # Return the clause:
            yield (mode, clause)


clauses_added = []
clauses_deleted = []

for mode, clause in tqdm.tqdm(parse_binary_drat_mmap("proof.drat")):
    if mode == "a":
        clauses_added.append(clause)
    elif mode == "d":
        clauses_deleted.append(clause)
    else:
        raise ValueError(f"Bad clause mode '{mode}'")

print(f"Parsed {len(clauses_added)} added clauses from DRAT file")
print(f"Parsed {len(clauses_deleted)} deleted clauses from DRAT file")