From 61bd78755f11ccbbdd06e53f6fcb0248720c880c Mon Sep 17 00:00:00 2001
From: Evan Hubinger <evanjhub@gmail.com>
Date: Fri, 3 Nov 2023 22:59:44 -0700
Subject: [PATCH] Improve profiling

---
 .gitignore                 |  5 ++++-
 Makefile                   | 29 ++++++++++++++++++++---------
 coconut/compiler/header.py |  2 ++
 coconut/compiler/util.py   |  2 +-
 coconut/requirements.py    |  1 +
 5 files changed, 28 insertions(+), 11 deletions(-)

diff --git a/.gitignore b/.gitignore
index 26f176d2b..96d716fb7 100644
--- a/.gitignore
+++ b/.gitignore
@@ -138,7 +138,10 @@ pyprover/
 bbopt/
 coconut-prelude/
 index.rst
-vprof.json
 /coconut/icoconut/coconut/
 __coconut_cache__/
+
+# Profiling
+vprof.json
 profile.svg
+profile.speedscope
diff --git a/Makefile b/Makefile
index 1128c283b..35da6ac15 100644
--- a/Makefile
+++ b/Makefile
@@ -25,27 +25,27 @@ dev-py3: clean setup-py3
 
 .PHONY: setup
 setup:
-	python -m ensurepip
+	-python -m ensurepip
 	python -m pip install --upgrade setuptools wheel pip pytest_remotedata cython
 
 .PHONY: setup-py2
 setup-py2:
-	python2 -m ensurepip
+	-python2 -m ensurepip
 	python2 -m pip install --upgrade "setuptools<58" wheel pip pytest_remotedata cython
 
 .PHONY: setup-py3
 setup-py3:
-	python3 -m ensurepip
+	-python3 -m ensurepip
 	python3 -m pip install --upgrade setuptools wheel pip pytest_remotedata cython
 
 .PHONY: setup-pypy
 setup-pypy:
-	pypy -m ensurepip
+	-pypy -m ensurepip
 	pypy -m pip install --upgrade "setuptools<58" wheel pip pytest_remotedata
 
 .PHONY: setup-pypy3
 setup-pypy3:
-	pypy3 -m ensurepip
+	-pypy3 -m ensurepip
 	pypy3 -m pip install --upgrade setuptools wheel pip pytest_remotedata
 
 .PHONY: install
@@ -337,10 +337,21 @@ profile-parser: export COCONUT_PURE_PYTHON=TRUE
 profile-parser:
 	coconut ./coconut/tests/src/cocotest/agnostic ./coconut/tests/dest/cocotest --force --jobs 0 --profile --verbose --stack-size 4096 --recursion-limit 4096 2>&1 | tee ./profile.log
 
-.PHONY: pyspy
-pyspy:
-	py-spy record -o profile.svg --native -- python -m coconut ./coconut/tests/src/cocotest/agnostic ./coconut/tests/dest/cocotest --force --jobs 0
-	open profile.svg
+.PHONY: open-speedscope
+open-speedscope:
+	npm install -g speedscope
+	speedscope ./profile.speedscope
+
+.PHONY: pyspy-purepy
+pyspy-purepy: export COCONUT_PURE_PYTHON=TRUE
+pyspy-purepy:
+	py-spy record -o profile.speedscope --format speedscope --subprocesses -- python -m coconut ./coconut/tests/src/cocotest/agnostic ./coconut/tests/dest/cocotest --force
+	open-speedscope
+
+.PHONY: pyspy-native
+pyspy-native:
+	py-spy record -o profile.speedscope --format speedscope --native -- python -m coconut ./coconut/tests/src/cocotest/agnostic ./coconut/tests/dest/cocotest --force --jobs 0
+	open-speedscope
 
 .PHONY: vprof-time
 vprof-time:
diff --git a/coconut/compiler/header.py b/coconut/compiler/header.py
index 3910880b7..5b1a68686 100644
--- a/coconut/compiler/header.py
+++ b/coconut/compiler/header.py
@@ -44,6 +44,7 @@
     univ_open,
     get_target_info,
     assert_remove_prefix,
+    memoize,
 )
 from coconut.compiler.util import (
     split_comment,
@@ -96,6 +97,7 @@ def minify_header(compiled):
 template_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "templates")
 
 
+@memoize()
 def get_template(template):
     """Read the given template file."""
     with univ_open(os.path.join(template_dir, template) + template_ext, "r") as template_file:
diff --git a/coconut/compiler/util.py b/coconut/compiler/util.py
index 9db8bf782..1dfbe7d34 100644
--- a/coconut/compiler/util.py
+++ b/coconut/compiler/util.py
@@ -514,7 +514,7 @@ def get_pyparsing_cache():
     else:  # on pyparsing we have to do this
         try:
             # this is sketchy, so errors should only be complained
-            return get_func_closure(packrat_cache.get.__func__)["cache"]
+            return get_func_closure(packrat_cache.set.__func__)["cache"]
         except Exception as err:
             complain(err)
             return {}
diff --git a/coconut/requirements.py b/coconut/requirements.py
index 55c293471..c2e9668a0 100644
--- a/coconut/requirements.py
+++ b/coconut/requirements.py
@@ -247,6 +247,7 @@ def everything_in(req_dict):
 
 extras["dev"] = uniqueify_all(
     everything_in(extras),
+    get_reqs("purepython"),
     get_reqs("dev"),
 )