Optimize match def defaults

evhub · Mar 1, 2024 · d0acf3d · d0acf3d
1 parent 5a331c4
commit d0acf3d
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 29 deletions.
diff --git a/coconut/compiler/grammar.py b/coconut/compiler/grammar.py
@@ -953,6 +953,7 @@ class Grammar(object):
         )
 
         atom_item = Forward()
+        const_atom = Forward()
         expr = Forward()
         star_expr = Forward()
         dubstar_expr = Forward()
@@ -1161,13 +1162,17 @@ class Grammar(object):
                 )
             )
         )
+        match_arg_default = Group(
+            const_atom("const")
+            | test("expr")
+        )
         match_args_list = Group(Optional(
             tokenlist(
                 Group(
                     (star | dubstar) + match
                     | star  # not star_sep because pattern-matching can handle star separators on any Python version
                     | slash  # not slash_sep as above
-                    | match + Optional(equals.suppress() + test)
+                    | match + Optional(equals.suppress() + match_arg_default)
                 ),
                 comma,
             )
@@ -1292,19 +1297,24 @@ class Grammar(object):
         lazy_items = Optional(tokenlist(test, comma))
         lazy_list = attach(lbanana.suppress() + lazy_items + rbanana.suppress(), lazy_list_handle)
 
-        known_atom = (
+        # for const_atom, value should be known at compile time
+        const_atom <<= (
             keyword_atom
-            | string_atom
             | num_atom
+            # typedef ellipsis must come before ellipsis
+            | typedef_ellipsis
+            | ellipsis
+        )
+        # for known_atom, type should be known at compile time
+        known_atom = (
+            const_atom
+            | string_atom
             | list_item
             | dict_literal
             | dict_comp
             | set_literal
             | set_letter_literal
             | lazy_list
-            # typedef ellipsis must come before ellipsis
-            | typedef_ellipsis
-            | ellipsis
         )
         atom = (
             # known_atom must come before name to properly parse string prefixes
@@ -2197,7 +2207,7 @@ class Grammar(object):
                 (
                     lparen.suppress()
                     + match
-                    + Optional(equals.suppress() + test)
+                    + Optional(equals.suppress() + match_arg_default)
                     + rparen.suppress()
                 ) | interior_name_match
             )

diff --git a/coconut/compiler/matching.py b/coconut/compiler/matching.py
@@ -307,38 +307,49 @@ def get_set_name_var(self, name):
         """Gets the var for checking whether a name should be set."""
         return match_set_name_var + "_" + name
 
-    def add_default_expr(self, assign_to, default_expr):
+    def add_default_expr(self, assign_to, default):
         """Add code that evaluates expr in the context of any names that have been matched so far
         and assigns the result to assign_to if assign_to is currently _coconut_sentinel."""
-        vars_var = self.get_temp_var()
-        add_names_code = []
-        for name in self.names:
-            add_names_code.append(
-                handle_indentation(
-                    """
+        default_expr, = default
+        if "const" in default:
+            self.add_def(handle_indentation("""
+if {assign_to} is _coconut_sentinel:
+    {assign_to} = {default_expr}
+            """.format(
+                assign_to=assign_to,
+                default_expr=default_expr,
+            )))
+        else:
+            internal_assert("expr" in default, "invalid match default tokens", default)
+            vars_var = self.get_temp_var()
+            add_names_code = []
+            for name in self.names:
+                add_names_code.append(
+                    handle_indentation(
+                        """
 if {set_name_var} is not _coconut_sentinel:
     {vars_var}["{name}"] = {set_name_var}
-                    """,
-                    add_newline=True,
-                ).format(
-                    set_name_var=self.get_set_name_var(name),
-                    vars_var=vars_var,
-                    name=name,
+                        """,
+                        add_newline=True,
+                    ).format(
+                        set_name_var=self.get_set_name_var(name),
+                        vars_var=vars_var,
+                        name=name,
+                    )
                 )
-            )
-        code = self.comp.reformat_post_deferred_code_proc(assign_to + " = " + default_expr)
-        self.add_def(handle_indentation("""
+            code = self.comp.reformat_post_deferred_code_proc(assign_to + " = " + default_expr)
+            self.add_def(handle_indentation("""
 if {assign_to} is _coconut_sentinel:
     {vars_var} = _coconut.globals().copy()
     {vars_var}.update(_coconut.locals())
     {add_names_code}_coconut_exec({code_str}, {vars_var})
     {assign_to} = {vars_var}["{assign_to}"]
-        """).format(
-            vars_var=vars_var,
-            add_names_code="".join(add_names_code),
-            assign_to=assign_to,
-            code_str=self.comp.wrap_str_of(code),
-        ))
+            """).format(
+                vars_var=vars_var,
+                add_names_code="".join(add_names_code),
+                assign_to=assign_to,
+                code_str=self.comp.wrap_str_of(code),
+            ))
 
     def register_name(self, name):
         """Register a new name at the current position."""