feat: Implement groups into parser, AST

Author: tkarabela

src/regex_automata/automata/nfa.py

@@ -20,6 +20,8 @@ def is_trivial(self) -> bool:
 class Transition:
     predicates: tuple[TransitionPredicate, ...]
     consume_char: bool = True
+    begin_group: int | None = None
+    end_group: int | None = None
     label: str = ""
 
     def matches(self, c_previous: int, c_next: int) -> bool:
@@ -33,12 +35,27 @@ def matches(self, c_previous: int, c_next: int) -> bool:
 
     @property
     def is_trivial_epsilon(self) -> bool:
-        return not self.consume_char and all(p.is_trivial for p in self.predicates)
+        return (
+            not self.consume_char and
+            self.begin_group is None and
+            self.end_group is None and
+            all(p.is_trivial for p in self.predicates)
+        )
 
     @classmethod
     def make_trivial_epsilon(cls) -> Self:
         return cls(predicates=(TransitionPredicate(),), consume_char=False, label="ε")
 
+    @classmethod
+    def make_begin_group(cls, number: int) -> Self:
+        return cls(predicates=(TransitionPredicate(),), consume_char=False, label=f"⟨begin group {number}⟩",
+                   begin_group=number)
+
+    @classmethod
+    def make_end_group(cls, number: int) -> Self:
+        return cls(predicates=(TransitionPredicate(),), consume_char=False, label=f"⟨end group {number}⟩",
+                   end_group=number)
+
 
 @dataclass
 class NFA:
@@ -48,18 +65,18 @@ class NFA:
     """
     states: list[int]
     initial_state: int
-    final_states: list[int]
+    final_states: set[int]
     transitions: dict[int, dict[Transition, set[int]]]
 
     def copy(self) -> "NFA":
         return deepcopy(self)
 
-    def renumber_states(self, x0: int) -> "NFA":
+    def renumber_states(self, x0: int = 0) -> "NFA":
         f = dict(zip(self.states, count(x0)))
         return NFA(
             states=[f[x] for x in self.states],
             initial_state=f[self.initial_state],
-            final_states=[f[x] for x in self.final_states],
+            final_states={f[x] for x in self.final_states},
             transitions={
                 f[x]: {p: {f[y] for y in ys} for p, ys in d.items()}
                 for x, d in self.transitions.items()
@@ -133,6 +150,6 @@ def get_trivial_epsilon_free_nfa(self) -> "NFA":
         return NFA(
             states=list(sorted(reachable_states)),
             initial_state=initial_state,
-            final_states=list(sorted(final_states)),
+            final_states=final_states,
             transitions=transitions,
-        )
+        ).renumber_states()

src/regex_automata/parser/ast.py

@@ -105,3 +105,15 @@ def iter_children(self) -> Iterator["AstNode"]:
 class AstEmpty(AstNode):
     def get_label(self) -> str:
         return "ε"
+
+
+@dataclass
+class AstGroup(AstNode):
+    number: int
+    u: AstNode
+
+    def get_label(self) -> str:
+        return f"group {self.number}"
+
+    def iter_children(self) -> Iterator["AstNode"]:
+        yield self.u

src/regex_automata/parser/ast_processor.py

@@ -1,14 +1,15 @@
 
 from .ast import AstNode, AstEmpty, AstConcatenation, AstUnion, AstRepetition, AstCharacterSet, AstIteration, \
-    AstBoundaryAssertion
+    AstBoundaryAssertion, AstGroup
 
 
 class ASTProcessor:
     def __init__(self, raw_ast: AstNode) -> None:
         self.raw_ast = raw_ast
 
     def get_processed_ast(self) -> AstNode:
-        return self.convert(self.raw_ast)
+        ast = self.convert(self.raw_ast)
+        return AstGroup(0, ast)
 
     def convert(self, node: AstNode) -> AstNode:
         match node:
@@ -26,6 +27,8 @@ def convert(self, node: AstNode) -> AstNode:
                 return self.convert_AstConcatenation(node)
             case AstBoundaryAssertion():
                 return self.convert_AstBoundaryAssertion(node)
+            case AstGroup():
+                return self.convert_AstGroup(node)
             case _:
                 return node
 
@@ -88,6 +91,9 @@ def convert_AstConcatenation(self, node: AstConcatenation) -> AstNode:
     def convert_AstBoundaryAssertion(self, node: AstBoundaryAssertion) -> AstNode:
         return node
 
+    def convert_AstGroup(self, node: AstGroup) -> AstNode:
+        return AstGroup(node.number, self.convert(node.u))
+
     @staticmethod
     def iterated_concatenation(node: AstNode, n: int) -> AstNode:
         # 0 -> AstEmpty == ""

src/regex_automata/parser/parser.py

@@ -2,7 +2,8 @@
 from typing import Type, TypeVar, NoReturn, ParamSpec, Callable
 
 from .tokens import Token, LPar, RPar, Repetition, Pipe, CharacterSet, BoundaryAssertion
-from .ast import AstNode, AstUnion, AstRepetition, AstCharacterSet, AstConcatenation, AstEmpty, AstBoundaryAssertion
+from .ast import AstNode, AstUnion, AstRepetition, AstCharacterSet, AstConcatenation, AstEmpty, AstBoundaryAssertion, \
+    AstGroup
 from ..errors import ParserError
 
 
@@ -24,6 +25,7 @@ def __init__(self, tokens: list[Token]):
         self.tokens = list(tokens)
         self.pos = -1
         self.string_pos = -1
+        self.group_number = 1
 
     def read(self, cls: Type[TToken]) -> TToken:
         self.pos += 1
@@ -54,6 +56,11 @@ def parse(self) -> AstNode:
             self.error("unread input remaining (expected end of input)")
         return root
 
+    def make_group(self, u: AstNode) -> AstGroup:
+        i = self.group_number
+        self.group_number += 1
+        return AstGroup(i, u)
+
     @rule
     def p1(self) -> AstNode:
         """
@@ -232,7 +239,7 @@ def p10(self) -> AstNode:
         # rpar
         _ = self.read(RPar)
 
-        return E
+        return self.make_group(E)
 
     @rule
     def p11(self) -> AstNode:

src/regex_automata/regex/nfa_builder.py

@@ -1,6 +1,6 @@
 from ..automata.rangeset import RangeSet, WORD_RANGESET, NONWORD_RANGESET
 from ..parser.ast import AstNode, AstCharacterSet, AstConcatenation, AstUnion, AstEmpty, AstIteration, \
-    AstBoundaryAssertion
+    AstBoundaryAssertion, AstGroup
 from ..automata.nfa import NFA, Transition, TransitionPredicate
 from ..parser.tokens import BoundaryAssertionSemantic
 
@@ -30,22 +30,24 @@ def convert(self, node: AstNode) -> NFA:
                 return self.convert_AstConcatenation(node)
             case AstBoundaryAssertion():
                 return self.covert_AstBoundaryAssertion(node)
+            case AstGroup():
+                return self.convert_AstGroup(node)
             case _:
                 raise NotImplementedError(f"Cannot convert node {node!r}")
 
     def convert_AstEmpty(self, _: AstEmpty) -> NFA:
         return NFA(
             states=[0],
             initial_state=0,
-            final_states=[0],
+            final_states={0},
             transitions={}
         )
 
     def convert_AstCharacter(self, node: AstCharacterSet) -> NFA:
         return NFA(
             states=[0, 1],
             initial_state=0,
-            final_states=[1],
+            final_states={1},
             transitions={0: {Transition(predicates=(TransitionPredicate(next=node.rs),), label=node.label): {1}}}
         )
 
@@ -55,7 +57,7 @@ def convert_AstIteration(self, node: AstIteration) -> NFA:
 
         for x in nfa.final_states:
             nfa.transitions.setdefault(x, {}).setdefault(Transition.make_trivial_epsilon(), set()).add(nfa.initial_state)
-        nfa.final_states = list(sorted(nfa.trivial_epsilon_closure(set(nfa.final_states))))
+        nfa.final_states = nfa.trivial_epsilon_closure(set(nfa.final_states))
         return nfa
 
     def convert_AstUnion(self, node: AstUnion) -> NFA:
@@ -66,7 +68,7 @@ def convert_AstUnion(self, node: AstUnion) -> NFA:
 
         nfa = nfa_u.copy()
         nfa.states += nfa_v.states
-        nfa.final_states += nfa_v.final_states
+        nfa.final_states |= nfa_v.final_states
         nfa.transitions.update(nfa_v.transitions)
         new_initial_state = max(nfa.states) + 1
         nfa.states.append(new_initial_state)
@@ -144,6 +146,21 @@ def covert_AstBoundaryAssertion(self, node: AstBoundaryAssertion) -> NFA:
         return NFA(
             states=[0, 1],
             initial_state=0,
-            final_states=[1],
+            final_states={1},
             transitions={0: {transition: {1}}}
         )
+
+    def convert_AstGroup(self, node: AstGroup) -> NFA:
+        nfa_u = self.convert(node.u).renumber_states(1)
+
+        start_state = 0
+        final_state = max(nfa_u.states) + 1
+
+        nfa = nfa_u.copy()
+        nfa.initial_state = start_state
+        nfa.states += [start_state, final_state]
+        nfa.final_states = {final_state}
+        nfa.transitions[start_state] = {Transition.make_begin_group(node.number): {nfa_u.initial_state}}
+        for s in nfa_u.final_states:
+            nfa.transitions.setdefault(s, {}).setdefault(Transition.make_end_group(node.number), set()).add(final_state)
+        return nfa

tests/test_parser.py

@@ -2,7 +2,7 @@
 
 from regex_automata.automata.rangeset import RangeSet
 from regex_automata.errors import TokenizerError
-from regex_automata.parser.ast import AstUnion, AstCharacterSet, AstConcatenation
+from regex_automata.parser.ast import AstUnion, AstCharacterSet, AstConcatenation, AstGroup
 from regex_automata.parser.parser import Parser
 from regex_automata.parser.tokenizer import Tokenizer
 
@@ -31,9 +31,12 @@ def test_parse_tree_union_parens():
     assert Parser(tokens).parse() == AstUnion(
         AstConcatenation(_ast_character_set("a"), _ast_character_set("b")),
         AstUnion(
-            AstUnion(
-                AstConcatenation(_ast_character_set("c"), _ast_character_set("d")),
-                AstConcatenation(_ast_character_set("e"), _ast_character_set("f")),
+            AstGroup(
+                1,
+                AstUnion(
+                    AstConcatenation(_ast_character_set("c"), _ast_character_set("d")),
+                    AstConcatenation(_ast_character_set("e"), _ast_character_set("f")),
+                )
             ),
             AstConcatenation(_ast_character_set("g"), _ast_character_set("h")),
         )