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

Rules with several alternatives now switch on the next token's type and
only try the alternatives that may start with it. Loops stop as soon as
the next token cannot start another iteration. Both use FIRST sets from
a new analysis, which also replaces the old first_sets.py calculator.

Co-authored-by: Pablo Galindo <pablogsal@gmail.com>

Author: lysnikolaou

Lib/test/test_peg_generator/test_c_generator.py

@@ -1,4 +1,5 @@
 import io
+import token
 import unittest
 from pathlib import PurePath
 from unittest import mock
@@ -27,6 +28,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
@@ -286,6 +291,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,
+            },
+        )
+
 
 if __name__ == "__main__":
     unittest.main()

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.