Fix perf improvements

CONTRIBUTING.md

@@ -155,20 +155,23 @@ After you've tested your changes locally, you'll want to add more permanent test
 1. Preparation:
     1. Run `make check-reqs` and update dependencies as necessary
     2. Run `sudo make format`
-    3. Make sure `make test`, `make test-pyright`, and `make test-easter-eggs` are passing
-    4. Ensure that `coconut --watch` can successfully compile files when they're modified
-    5. Check changes in [`compiled-cocotest`](https://github.com/evhub/compiled-cocotest), [`pyprover`](https://github.com/evhub/pyprover), and [`coconut-prelude`](https://github.com/evhub/coconut-prelude)
-    6. Check [Codebeat](https://codebeat.co/a/evhub/projects) and [LGTM](https://lgtm.com/dashboard) for `coconut` and `compiled-cocotest`
-    7. Make sure [`coconut-develop`](https://pypi.python.org/pypi/coconut-develop) package looks good
-    8. Run `make docs` and ensure local documentation looks good
-    9. Make sure all of the following are passing:
+    3. Run manual tests:
+       1. Temporarily set `test_computation_graph_pickling = True` in `constants.py` and run `make test`
+       2. Run `make docs` and ensure local documentation looks good
+       3. Ensure that `coconut --watch` can successfully compile files when they're modified
+       4. Make sure `make test-pyright` and `make test-easter-eggs` are passing
+    4. Check all the tooling:
+       1. Make sure [`coconut-develop`](https://pypi.python.org/pypi/coconut-develop) package looks good
+       2. Check changes in [`compiled-cocotest`](https://github.com/evhub/compiled-cocotest), [`pyprover`](https://github.com/evhub/pyprover), and [`coconut-prelude`](https://github.com/evhub/coconut-prelude)
+       3. Check [Codebeat](https://codebeat.co/a/evhub/projects) and [LGTM](https://lgtm.com/dashboard) for `coconut` and `compiled-cocotest`
+    5.  Make sure all of the following are passing:
         1.  [Github Actions](https://github.com/evhub/coconut/actions)
         2.  [AppVeyor](https://ci.appveyor.com/project/evhub/coconut)
         3.  [readthedocs](https://readthedocs.org/projects/coconut/builds/)
-    10. Make sure [develop documentation](http://coconut.readthedocs.io/en/develop/) looks good
-    11. Turn off `develop` in `root.py`
-    12. Set `root.py` to new version number
-    13. If major release, set `root.py` to new version name
+    6.  Make sure [develop documentation](http://coconut.readthedocs.io/en/develop/) looks good
+    7.  Turn off `develop` in `root.py`
+    8.  Set `root.py` to new version number
+    9.  If major release, set `root.py` to new version name
 
 2. Pull Request:
     1. Move unresolved issues to new milestone

coconut/compiler/compiler.py

@@ -618,7 +618,8 @@ def reset(self, keep_state=False, filename=None):
         self.add_code_before_ignore_names = {}
         self.remaining_original = None
         self.shown_warnings = set()
-        self.computation_graph_caches = defaultdict(staledict)
+        if not keep_state:
+            self.computation_graph_caches = defaultdict(staledict)
 
     @contextmanager
     def inner_environment(self, ln=None):
@@ -2396,28 +2397,27 @@ def split_docstring(self, block):
                 return first_line, rest_of_lines
         return None, block
 
-    def tre_return_grammar(self, func_name, func_args, func_store, mock_var=None):
-        """Generate grammar element that matches a string which is just a TRE return statement."""
-        def tre_return_handle(loc, tokens):
-            args = ", ".join(tokens)
+    def tre_return_handle(self, func_name, func_args, func_store, mock_var, loc, tokens):
+        """Handler for tre_return_grammar."""
+        args = ", ".join(tokens)
 
-            # we have to use func_name not func_store here since we use
-            #  this when we fail to verify that func_name is func_store
-            if self.no_tco:
-                tco_recurse = "return " + func_name + "(" + args + ")"
-            else:
-                tco_recurse = "return _coconut_tail_call(" + func_name + (", " + args if args else "") + ")"
+        # we have to use func_name not func_store here since we use
+        #  this when we fail to verify that func_name is func_store
+        if self.no_tco:
+            tco_recurse = "return " + func_name + "(" + args + ")"
+        else:
+            tco_recurse = "return _coconut_tail_call(" + func_name + (", " + args if args else "") + ")"
 
-            if not func_args or func_args == args:
-                tre_recurse = "continue"
-            elif mock_var is None:
-                tre_recurse = tuple_str_of_str(func_args) + " = " + tuple_str_of_str(args) + "\ncontinue"
-            else:
-                tre_recurse = tuple_str_of_str(func_args) + " = " + mock_var + "(" + args + ")" + "\ncontinue"
+        if not func_args or func_args == args:
+            tre_recurse = "continue"
+        elif mock_var is None:
+            tre_recurse = tuple_str_of_str(func_args) + " = " + tuple_str_of_str(args) + "\ncontinue"
+        else:
+            tre_recurse = tuple_str_of_str(func_args) + " = " + mock_var + "(" + args + ")" + "\ncontinue"
 
-            tre_check_var = self.get_temp_var("tre_check", loc)
-            return handle_indentation(
-                """
+        tre_check_var = self.get_temp_var("tre_check", loc)
+        return handle_indentation(
+            """
 try:
     {tre_check_var} = {func_name} is {func_store} {type_ignore}
 except _coconut.NameError:
@@ -2426,21 +2426,23 @@ def tre_return_handle(loc, tokens):
     {tre_recurse}
 else:
     {tco_recurse}
-                """,
-                add_newline=True,
-            ).format(
-                tre_check_var=tre_check_var,
-                func_name=func_name,
-                func_store=func_store,
-                tre_recurse=tre_recurse,
-                tco_recurse=tco_recurse,
-                type_ignore=self.type_ignore_comment(),
-            )
+            """,
+            add_newline=True,
+        ).format(
+            tre_check_var=tre_check_var,
+            func_name=func_name,
+            func_store=func_store,
+            tre_recurse=tre_recurse,
+            tco_recurse=tco_recurse,
+            type_ignore=self.type_ignore_comment(),
+        )
+
+    def tre_return_grammar(self, func_name, func_args, func_store, mock_var=None):
+        """Generate grammar element that matches a string which is just a TRE return statement."""
         self.tre_func_name <<= base_keyword(func_name).suppress()
         return StartOfStrGrammar(attach(
             self.tre_return_base,
-            tre_return_handle,
-            greedy=True,
+            partial(self.tre_return_handle, func_name, func_args, func_store, mock_var),
         ))
 
     def detect_is_gen(self, raw_lines):