re2dfa.py

from abc import ABC, abstractmethod
from collections import defaultdict
import functools
import itertools

INPUT_FILE = 'regex.in'
OUTPUT_FILE = 'dfa.out'


# Based on http://matt.might.net/articles/parsing-regex-with-recursive-descent/
class Parser:
    """Parses a string as a regular expression and constructs an abstract
    syntax tree."""

    def __init__(self, input):
        self.input = input
        self.input_len = len(input)
        self.cursor = 0
        self.id = 0

    @property
    def position(self):
        return self.cursor + 1

    @property
    def symbol(self):
        return self.input[self.cursor]

    def expect(self, expected):
        if self.symbol == expected:
            self.cursor += 1
        else:
            raise Exception(f"Expected '{expected}' at {self.position}, " +
                            "got '{self.symbol}' instead")

    def next(self):
        """Consumes the next symbol from the output and returns it."""
        sym = self.symbol
        self.expect(sym)
        return sym

    def more(self):
        """Returns true if there still are some characters in the input."""
        return self.cursor < self.input_len

    def parse(self):
        root = self.expression()
        if self.cursor != self.input_len:
            raise Exception("More input left after expression")
        return root

    def expression(self):
        term = self.term()

        if self.more() and self.symbol == '|':
            self.expect('|')
            rhs = self.expression()
            return OrNode(term, rhs)

        return term

    def term(self):
        factor = self.factor()

        while self.more() and self.symbol != ')' and self.symbol != '|':
            rhs = self.factor()
            factor = ConcatNode(factor, rhs)

        return factor

    def factor(self):
        base = self.base()

        while self.more() and self.symbol == '*':
            self.expect('*')
            base = StarNode(base)

        return base

    def base(self):
        sym = self.symbol
        if sym == '(':
            self.expect('(')
            expr = self.expression()
            self.expect(')')
            return expr
        else:
            self.id += 1
            return SymbolNode(self.id, self.next())


# Based on:
# http://staff.cs.upt.ro/~chirila/teaching/upt/mse11-cd/slides/cd0309.pdf
class Node(ABC):
    @property
    @abstractmethod
    def nullable(self) -> bool:
        """Returns true if lambda is contained in the language
        represented by this node's subtree.
        """

    @property
    @abstractmethod
    def first_pos(self) -> frozenset:
        """Returns the set of positions in this subtree which can begin a word
        accepted by this subexpression's language.
        """

    @property
    @abstractmethod
    def last_pos(self) -> frozenset:
        """Returns the set of positions in this subtree which can end a word
        accepted by this subexpression's language.
        """

    @property
    @abstractmethod
    def alphabet(self) -> frozenset:
        """Returns the set of symbols which appear in this subtree."""

    def compute_follow_pos(self, follow_pos) -> None:
        """Computes the follow positions of this subtree recursively."""

    def compute_follow_union(self, follow_pos, union, state, symbol) -> None:
        """Computes the union of possible follow positions for all the states
        in a superstate, through the given symbol."""


class SymbolNode(Node):
    """Leaf node of the AST, generated by a literal symbol in the input regex,
    corresponding to an important position in the final DFA.
    """

    def __init__(self, position, symbol):
        assert len(symbol) == 1, "Symbol should be only 1 character"

        self.position = position
        self.symbol = symbol

    def is_lambda(self):
        return self.symbol == '.'

    @property
    def nullable(self):
        return self.is_lambda()

    @property
    def first_pos(self):
        return frozenset() if self.is_lambda() else frozenset((self.position,))

    @property
    def last_pos(self):
        return frozenset() if self.is_lambda() else frozenset((self.position,))

    @property
    def alphabet(self):
        return frozenset() if self.is_lambda() else frozenset((self.symbol,))

    def compute_follow_union(self, follow_pos, union, state, symbol):
        if self.position in state and self.symbol == symbol:
            union |= follow_pos[self.position]

    def __repr__(self):
        return self.symbol


class OrNode(Node):
    """Node representing the disjunction between its left and right child."""

    def __init__(self, left, right):
        self.left = left
        self.right = right

    @property
    def nullable(self):
        return self.left.nullable or self.right.nullable

    @property
    def first_pos(self):
        return self.left.first_pos | self.right.first_pos

    @property
    def last_pos(self):
        return self.left.last_pos | self.right.last_pos

    @property
    def alphabet(self):
        return self.left.alphabet | self.right.alphabet

    def compute_follow_pos(self, follow_pos):
        self.left.compute_follow_pos(follow_pos)
        self.right.compute_follow_pos(follow_pos)

    def compute_follow_union(self, follow_pos, union, state, symbol):
        self.left.compute_follow_union(follow_pos, union, state, symbol)
        self.right.compute_follow_union(follow_pos, union, state, symbol)

    def __repr__(self):
        return f'({self.left}|{self.right})'


class ConcatNode(Node):
    """Concatenates the expression represented by the left subtree with the one
    represented by the right subtree.
    """

    def __init__(self, left, right):
        self.left = left
        self.right = right

    @property
    def nullable(self):
        return self.left.nullable and self.right.nullable

    @property
    def first_pos(self):
        if self.left.nullable:
            return self.left.first_pos | self.right.first_pos
        else:
            return self.left.first_pos

    @property
    def last_pos(self):
        if self.right.nullable:
            return self.right.last_pos | self.left.last_pos
        else:
            return self.right.last_pos

    @property
    def alphabet(self):
        return self.left.alphabet | self.right.alphabet

    def compute_follow_pos(self, follow_pos):
        """When concatenating two nodes, the last positions"""
        for position in self.left.last_pos:
            follow_pos[position] |= self.right.first_pos

        self.left.compute_follow_pos(follow_pos)
        self.right.compute_follow_pos(follow_pos)

    def compute_follow_union(self, follow_pos, union, state, symbol):
        self.left.compute_follow_union(follow_pos, union, state, symbol)
        self.right.compute_follow_union(follow_pos, union, state, symbol)

    def __repr__(self):
        return f'({self.left}{self.right})'


class StarNode(Node):
    """Node applying the Kleene star to its child node."""

    def __init__(self, child):
        self.child = child

    @property
    def nullable(self):
        """Starring always also contains the empty word."""
        return True

    @property
    def first_pos(self):
        return self.child.first_pos

    @property
    def last_pos(self):
        return self.child.last_pos

    @property
    def alphabet(self):
        return self.child.alphabet

    def compute_follow_pos(self, follow_pos):
        # Since a star can contain its child concatenated with another copy
        # of its child, the last position of its child can be followed by the
        # first position of the other copy of the child.
        for position in self.last_pos:
            follow_pos[position] |= self.first_pos

        self.child.compute_follow_pos(follow_pos)

    def compute_follow_union(self, follow_pos, union, state, symbol):
        self.child.compute_follow_union(follow_pos, union, state, symbol)

    def __repr__(self):
        return f'({self.child}*)'


class StateMapper:
    """Dictionary-like class which maps frozen sets of states to
    new states with numeric IDs."""

    def __init__(self):
        counter = itertools.count(1)
        self.mappings = defaultdict(lambda: next(counter))

    def __getitem__(self, index):
        return self.mappings[index]

    def final_states(self):
        return [
            state_id for state, state_id
            in self.mappings.items()
            if -1 in state
        ]

    def value_list(self):
        return list(self.mappings.values())


def regex_to_dfa(root):
    """Builds a DFA which accepts the same language as the regular expression
    stored in the abstract syntax tree passed as an argument to this function.

    Returns a 5-tuple representing the DFA:
    - the set of states
    - the alphabet
    - the transition function
    - the initial state
    - the set of final states.
    """

    root = ConcatNode(root, SymbolNode(-1, '#'))
    follow_pos = defaultdict(lambda: set())
    root.compute_follow_pos(follow_pos)

    alphabet = root.left.alphabet

    visited = set()
    queue = [root.first_pos]

    transitions = {}

    state_mapper = StateMapper()

    while queue:
        state = queue.pop()
        state_id = state_mapper[state]
        visited.add(state_id)

        for symbol in alphabet:
            union = set()
            root.compute_follow_union(follow_pos, union, state, symbol)
            union = frozenset(union)
            union_id = state_mapper[union]

            if union_id not in visited:
                queue.append(union)

            transitions[(state_id, symbol)] = union_id

    states = state_mapper.value_list()
    alphabet = list(alphabet)
    initial_state = 1
    final_states = state_mapper.final_states()

    return states, alphabet, transitions, initial_state, final_states


def read_regex(file):
    """Reads a regular expression from a given files-like object,
    presumed to be in text mode.

    Returns the AST of the parsed regex.
    """
    line = file.readline().strip()

    return Parser(line).parse()


def print_dfa(dfa, file):
    """Prints the 5-tuple representing a DFA to a file-like object.
    """

    states, alphabet, transitions, initial_state, final_states = dfa

    fprint = functools.partial(print, file=file)

    fprint(len(states))
    fprint(*states, sep=' ')

    fprint(len(alphabet))
    fprint(*alphabet, sep=' ')

    fprint(initial_state)

    fprint(len(final_states))
    fprint(*final_states, sep=' ')

    fprint(len(transitions))
    for (start, symbol), end in transitions.items():
        fprint(start, symbol, end, sep=' ')


def main():
    with open(INPUT_FILE, 'r') as f:
        re = read_regex(f)

    dfa = regex_to_dfa(re)

    with open(OUTPUT_FILE, 'w') as f:
        print_dfa(dfa, f)


if __name__ == '__main__':
    main()