Three Dollar Quill

"""Prototype implementation of NFA-based regex engine.

See https://swtch.com/~rsc/regexp/regexp1.html
"""

# My old nemesis
# pylint: disable=too-few-public-methods

import enum
import typing

import typing_extensions

# I think pylint is getting a little big for its britches when it tells us not
# to follow the usage given by the official documentation for the typing module
Index = typing.NewType('Index', int)  # pylint: disable=invalid-name
Grapheme = typing.NewType('Grapheme', str)  # pylint: disable=invalid-name


class Tail(typing.NamedTuple):
    """A "pointer" to an out-field of a State in an arena."""

    source: Index
    field: str


_T = typing.TypeVar('_T')


class State(typing_extensions.Protocol):
    """States have tails, and the ability to assign to tails."""

    # I really hope there's a plugin or something for dealing with Protocols.
    # pylint: disable=no-self-use, unused-argument

    def tails(self) -> typing.Iterator[Tail]:
        """Return the state's tails."""
        # I don't know if protocols NEED-need this, I'm just following the doc.
        ...  # pylint: disable=pointless-statement

    # There's probably a more correct interface to this...
    def _replace(self: _T, **kwargs) -> _T:
        """Point the state's tail at an index."""
        # I don't know if protocols NEED-need this, I'm just following the doc.
        ...  # pylint: disable=pointless-statement


class GraphemeMatch(typing.NamedTuple):
    """Match a single grapheme."""

    # When I extend it, I end up noticing that tuple has an obnoxiously big API
    # surface. So many common words I can't use because all methods go in the
    # same namespace...
    index_: Index
    grapheme: Grapheme
    out: typing.Optional[Index] = None

    def tails(self) -> typing.Iterator[Tail]:
        """Yield the only followup state."""
        yield Tail(self.index_, 'out')


class Split(typing.NamedTuple):
    """Split NFA evaluation between two possible states."""

    index_: Index
    out: typing.Optional[Index] = None
    out1: typing.Optional[Index] = None

    def tails(self) -> typing.Iterator[Tail]:
        """Yield both followup states."""
        yield Tail(self.index_, 'out')
        yield Tail(self.index_, 'out1')


class Match(typing.NamedTuple):
    """Terminate NFA evaluation."""

    index_: Index

    @staticmethod
    def tails() -> typing.Iterator[Tail]:
        """Yield nothing, there is no followup state."""
        yield from iter(())


class Fragment(typing.NamedTuple):
    """NFA fragment."""

    start: Index
    tail_list: typing.Tuple[Tail, ...]


class NFA(typing.NamedTuple):
    """Completed NFA of a regex."""

    start_index: Index
    arena: typing.Tuple[State, ...]

    def _addstate(self, indices: typing.Set[Index], index: Index):
        state = self.arena[index]
        if isinstance(state, Split):
            self._addstate(indices, state.out)
            self._addstate(indices, state.out1)
            return
        indices.add(index)

    def match(self, stream: typing.Iterable[Grapheme]) -> bool:
        """Return whether given grapheme stream is accepted by the regex."""
        indices: typing.Set[Index] = set()
        self._addstate(indices, self.start_index)
        for grapheme in stream:
            new_indices: typing.Set[Index] = set()
            for index in indices:
                state = self.arena[index]
                if isinstance(
                        state, GraphemeMatch) and state.grapheme == grapheme:
                    self._addstate(new_indices, state.out)
            indices = new_indices
        final_index = len(self.arena) - 1
        return final_index in indices


def patch(
        arena: typing.List[State],
        tail_list: typing.Sequence[Tail],
        index: Index):
    """Connect all tails to the given index."""
    for tail in tail_list:
        arena[tail.source] = arena[tail.source]._replace(**{tail.field: index})


class Token(typing_extensions.Protocol):
    """Protocol for postfix tokens."""

    # I really hope there's a plugin or something for dealing with Protocols.
    # pylint: disable=no-self-use, unused-argument

    def fragment_args(
            self, stack: typing.List[Fragment]) -> typing.Iterable[Fragment]:
        """Return the fragment args required by the fragment() call."""
        # I don't know if protocols NEED-need this, I'm just following the doc.
        ...  # pylint: disable=pointless-statement

    def fragment(
            self,
            arena: typing.List[State],
            fragment_args: typing.Iterable[Fragment]) -> Fragment:
        """Create a new fragment."""
        # I don't know if protocols NEED-need this, I'm just following the doc.
        ...  # pylint: disable=pointless-statement


class Literal(typing.NamedTuple):
    """Postfix token for matching a grapheme."""

    grapheme: Grapheme

    # pylint: disable=unused-argument
    @staticmethod
    def fragment_args(
            stack: typing.List[Fragment]) -> typing.Iterable[Fragment]:
        """No need to manipulate the stack."""
        return ()

    def fragment(
            self,
            arena: typing.List[State],
            fragment_args: typing.Iterable[Fragment]) -> Fragment:
        """Create a completely new Fragment."""
        index = push(arena, GraphemeMatch, self.grapheme)
        return Fragment(index, tuple(arena[index].tails()))


class Operator:
    """Postfix token for operators.

    This class may not be necessary.
    """

    def fragment_args(
            self, stack: typing.List[Fragment]) -> typing.Iterable[Fragment]:
        """Return the fragment_args required by the fragment() call."""
        raise NotImplementedError

    def fragment(
            self,
            arena: typing.List[State],
            fragment_args: typing.Iterable[Fragment]) -> Fragment:
        """Create a new fragment."""
        raise NotImplementedError


class Binary(Operator, enum.Enum):
    """Postfix token for binary operators."""

    CONCATENATE = enum.auto()
    ALTERNATE = enum.auto()

    def fragment_args(
            self, stack: typing.List[Fragment]) -> typing.Iterable[Fragment]:
        fragment_2 = stack.pop()
        fragment_1 = stack.pop()
        return fragment_1, fragment_2

    def fragment(
            self,
            arena: typing.List[State],
            fragment_args: typing.Iterable[Fragment]) -> Fragment:
        fragment_1, fragment_2 = fragment_args
        if self == Binary.CONCATENATE:
            patch(arena, fragment_1.tail_list, fragment_2.start)
            return Fragment(fragment_1.start, fragment_2.tail_list)
        assert self == Binary.ALTERNATE
        index = push(arena, Split, fragment_1.start, fragment_2.start)
        return Fragment(index, fragment_1.tail_list + fragment_2.tail_list)


class Unary(Operator, enum.Enum):
    """Postfix token for unary operators."""

    OPTIONAL = enum.auto()
    ANY_AMOUNT = enum.auto()
    ONE_PLUS = enum.auto()

    def fragment_args(
            self, stack: typing.List[Fragment]) -> typing.Iterable[Fragment]:
        return stack.pop(),

    def fragment(
            self,
            arena: typing.List[State],
            fragment_args: typing.Iterable[Fragment]):
        fragment, = fragment_args
        index = push(arena, Split, fragment.start)
        if self == Unary.OPTIONAL:
            return Fragment(
                index, fragment.tail_list + tuple(arena[index].tails()))
        elif self == Unary.ANY_AMOUNT:
            patch(arena, fragment.tail_list, index)
            return Fragment(index, (Tail(index, 'out1'),))
        assert self == Unary.ONE_PLUS
        patch(arena, fragment.tail_list, index)
        return Fragment(fragment.start, (Tail(index, 'out1'),))


def push(arena: typing.List[State], factory, *args, **kwargs):
    """Create a state using given factory, add to arena, return index."""
    index = Index(len(arena))
    state = factory(index, *args, **kwargs)
    arena.append(state)
    return index


def post_to_nfa(tokens: typing.Sequence[Token]):
    """Convert a list of postfix tokens to an NFA."""
    stack: typing.List[Fragment] = []
    arena: typing.List[State] = []
    for token in tokens:
        stack.append(token.fragment(arena, token.fragment_args(stack)))

    # pylint pointed out a bug on this next line. It's possible to hit it by
    # passing an empty string, which raises the question of how the C
    # implementation handles it. I don't know C well enough, so I'm just
    # staring at the for-loop and wondering "So, like, does that end if the
    # current character is null, or something?"
    fragment, = stack

    index = push(arena, Match)
    patch(arena, fragment.tail_list, index)
    return NFA(fragment.start, tuple(arena))