gh-153568: Skip PEG alternatives that cannot match the next token

Lib/test/test_peg_generator/test_c_generator.py

@@ -1,4 +1,5 @@
 import io
+import token
 import textwrap
 import unittest
 from pathlib import Path, PurePath
@@ -30,6 +31,10 @@ def emit_parser(self, parser):
         CParserEmitter(parser, output).emit()
         return output.getvalue()
 
+    def prepare_rules(self, source):
+        parser = self.make_generator(source).prepare(PurePath("example.gram"))
+        return {rule.signature.name: rule for rule in parser.rules}
+
     def test_rule_types_distinguish_implicit_and_explicit_void_pointer(self):
         generator = self.make_generator("""
             start: &implicit implicit explicit
@@ -289,6 +294,106 @@ def test_keywords_are_numbered_once(self):
         self.assertEqual(generator.keywords, {"if": 500, "else": 501, "pass": 502})
         self.assertEqual(generator.soft_keywords, {"soft"})
 
+    def test_first_set_dispatch(self):
+        rules = self.prepare_rules("""
+            start: NAME | '(' NAME | 'if' NAME | "match" NUMBER
+        """)
+        self.assertEqual(
+            rules["start"].first_set_dispatch,
+            (
+                frozenset({token.NAME}),
+                frozenset({token.LPAR}),
+                # Keyword types are numbered from 500.
+                frozenset({500}),
+                # Soft keywords are NAME tokens.
+                frozenset({token.NAME}),
+            ),
+        )
+
+    def test_first_set_dispatch_ignores_second_pass_alternatives(self):
+        rules = self.prepare_rules("""
+            start: NAME | invalid_start | NUMBER
+            single: NAME | invalid_start
+            invalid_start: STRING
+        """)
+        self.assertEqual(
+            rules["start"].first_set_dispatch,
+            (frozenset({token.NAME}), frozenset(), frozenset({token.NUMBER})),
+        )
+        self.assertIsNone(rules["single"].first_set_dispatch)
+
+    def test_no_first_set_dispatch(self):
+        cases = {
+            "single": "NAME",
+            "same_sets": "NAME '+' NAME | NAME",
+            "unknown": "~ NAME | [NUMBER]",
+        }
+        rules = self.prepare_rules(
+            "start: " + " | ".join(cases) + "\n"
+            + "".join(f"{name}: {rhs}\n" for name, rhs in cases.items())
+        )
+        for name in cases:
+            with self.subTest(name):
+                self.assertIsNone(rules[name].first_set_dispatch)
+
+    def test_first_set_dispatch_in_second_pass_rules(self):
+        rules = self.prepare_rules("""
+            start: NAME | invalid_start
+            invalid_start: (NAME | NUMBER) helper
+            helper: NAME | NUMBER
+        """)
+        self.assertEqual(
+            rules["helper"].first_set_dispatch,
+            (frozenset({token.NAME}), frozenset({token.NUMBER})),
+        )
+
+    def test_first_set_dispatch_emission(self):
+        generator = self.make_generator("""
+            start: disjoint | overlapping
+            disjoint: NAME | NUMBER
+            overlapping: NAME '+' NAME | NUMBER | ~ STRING
+        """)
+        source = self.emit_parser(generator.prepare(PurePath("example.gram")))
+        disjoint = source[source.index("\ndisjoint_rule(") : source.index("\noverlapping_rule(")]
+        self.assertIn(f"case {token.NAME}:\n                goto _first_set_alt_0;", disjoint)
+        self.assertNotIn("_first_set_mask", disjoint)
+        overlapping = source[source.index("\noverlapping_rule(") :]
+        # The alternative with an unknown first set is tried for every token.
+        self.assertIn(
+            f"case {token.NAME}:\n                _first_set_mask = 0x5ULL;", overlapping
+        )
+        self.assertIn("default:\n                _first_set_mask = 0x4ULL;", overlapping)
+
+    def test_loop_first_set_check(self):
+        rules = self.prepare_rules("""
+            start: element+ keyword* NAME* large*
+            element: NAME | NUMBER
+            keyword: 'if' NAME
+            large: tstring_end NAME
+            tstring_end: TSTRING_END
+        """)
+        checks = {}
+        for name, rule in rules.items():
+            if rule.signature.kind in {RuleKind.LOOP0, RuleKind.LOOP1}:
+                alt, = rule.alternatives
+                first = alt.calls[0]
+                if first.function == "_PyPegen_next_token_in":
+                    checks[str(alt.text)] = first.arguments
+                else:
+                    checks[str(alt.text)] = None
+        self.assertEqual(
+            checks,
+            {
+                "element": ("p", f"0x{1 << token.NAME | 1 << token.NUMBER:x}ULL", "0x0ULL"),
+                # Keyword types are numbered from 500.
+                "keyword": ("p", "0x0ULL", "0x1ULL"),
+                # Matching the token costs no more than checking it.
+                "NAME": None,
+                # The mask has no bit for a token type above 63.
+                "large": None,
+            },
+        )
+
 
 class TestIncludedCFiles(unittest.TestCase):
     def setUp(self):

Lib/test/test_peg_generator/test_first_sets.py

@@ -7,14 +7,18 @@
 with test_tools.imports_under_tool("peg_generator"):
     from pegen.grammar_parser import GeneratedParser as GrammarParser
     from pegen.testutil import parse_string
-    from pegen.first_sets import FirstSetCalculator
+    from pegen.grammar_analysis import FirstSetVisitor
     from pegen.grammar import Grammar
 
 
 class TestFirstSets(unittest.TestCase):
-    def calculate_first_sets(self, grammar_source: str) -> Dict[str, Set[str]]:
+    def calculate_first_sets(self, grammar_source: str) -> Dict[str, Set[str] | None]:
         grammar: Grammar = parse_string(grammar_source, GrammarParser)
-        return FirstSetCalculator(grammar.rules).calculate()
+        first_sets = FirstSetVisitor(grammar.rules).first_sets
+        return {
+            name: None if first is None else set(first)
+            for name, first in first_sets.items()
+        }
 
     def test_alternatives(self) -> None:
         grammar = """
@@ -178,10 +182,12 @@ def test_negative_lookahead(self) -> None:
         """
         self.assertEqual(
             self.calculate_first_sets(grammar),
+            # A negative lookahead doesn't narrow the first set, but its own
+            # first set is included for the effects it may have.
             {
                 "opt": {"'b'", "'a'", "'c'"},
-                "expr": {"'b'", "'c'"},
-                "start": {"'b'", "'c'"},
+                "expr": {"'a'", "'b'", "'c'"},
+                "start": {"'a'", "'b'", "'c'"},
             },
         )
 
@@ -213,7 +219,8 @@ def test_advance_left_recursion(self) -> None:
         """
         self.assertEqual(
             self.calculate_first_sets(grammar),
-            {"sign": {"'-'", ""}, "start": {"'-'", "NUMBER"}},
+            # sign may match whatever the next token is.
+            {"sign": None, "start": None},
         )
 
     def test_mutual_left_recursion(self) -> None:
@@ -226,21 +233,20 @@ def test_mutual_left_recursion(self) -> None:
             self.calculate_first_sets(grammar),
             {
                 "foo": {"'D'", "'B'"},
-                "bar": {"'D'"},
+                "bar": {"'D'", "'B'"},
                 "start": {"'D'", "'B'"},
             },
         )
 
     def test_nasty_left_recursion(self) -> None:
-        # TODO: Validate this
         grammar = """
         start: target '='
         target: maybe '+' | NAME
         maybe: maybe '-' | target
         """
         self.assertEqual(
             self.calculate_first_sets(grammar),
-            {"maybe": set(), "target": {"NAME"}, "start": {"NAME"}},
+            {"maybe": {"NAME"}, "target": {"NAME"}, "start": {"NAME"}},
         )
 
     def test_nullable_rule(self) -> None:
@@ -252,9 +258,9 @@ def test_nullable_rule(self) -> None:
         self.assertEqual(
             self.calculate_first_sets(grammar),
             {
-                "sign": {"", "'-'"},
+                "sign": None,
                 "thing": {"NUMBER"},
-                "start": {"NUMBER", "'-'"},
+                "start": None,
             },
         )
 
@@ -277,10 +283,60 @@ def test_multiple_nullable_rules(self) -> None:
         self.assertEqual(
             self.calculate_first_sets(grammar),
             {
-                "sign": {"", "'-'"},
-                "thing": {"'+'", ""},
-                "start": {"'+'", "'-'", "'*'"},
+                "sign": None,
+                "thing": None,
+                "start": None,
                 "other": {"'*'"},
                 "another": {"'/'"},
             },
         )
+
+    def test_keywords_and_exact_tokens(self) -> None:
+        grammar = """
+        start: 'if' NAME | '(' NAME ')'
+        """
+        self.assertEqual(self.calculate_first_sets(grammar), {"start": {"'if'", "'('"}})
+
+    def test_soft_keywords(self) -> None:
+        grammar = """
+        start: "match" NAME | SOFT_KEYWORD NUMBER
+        """
+        self.assertEqual(self.calculate_first_sets(grammar), {"start": {"NAME"}})
+
+    def test_leading_cut(self) -> None:
+        grammar = """
+        start: choice NEWLINE
+        choice: ~ NAME | NUMBER
+        """
+        self.assertEqual(
+            self.calculate_first_sets(grammar), {"choice": None, "start": None}
+        )
+
+    def test_forced(self) -> None:
+        grammar = """
+        start: choice NEWLINE
+        choice: &&'+' | NAME &&':' NUMBER
+        """
+        self.assertEqual(
+            self.calculate_first_sets(grammar), {"choice": None, "start": None}
+        )
+        grammar = """
+        start: NAME &&':' NUMBER
+        """
+        self.assertEqual(self.calculate_first_sets(grammar), {"start": {"NAME"}})
+
+    def test_negative_lookahead_effects(self) -> None:
+        grammar = """
+        start: !guard NAME
+        guard: NUMBER
+        """
+        self.assertEqual(
+            self.calculate_first_sets(grammar),
+            {"guard": {"NUMBER"}, "start": {"NAME", "NUMBER"}},
+        )
+
+    def test_positive_lookahead_alone(self) -> None:
+        grammar = """
+        start: &NUMBER
+        """
+        self.assertEqual(self.calculate_first_sets(grammar), {"start": {"NUMBER"}})

Misc/NEWS.d/next/Core_and_Builtins/2026-10-23-18-07-52.gh-issue-153568.Y6LTe5.rst

@@ -0,0 +1,3 @@
+Speed up the parser by computing the first set of each rule and skipping
+alternatives that cannot start with the next token based on those first
+sets.