Merge pull request #4057 from jtraglia/speed-up-make-pyspec

Speed up make pyspec command

setup.py

@@ -12,7 +12,7 @@
 import copy
 from collections import OrderedDict
 import json
-from functools import reduce
+from functools import lru_cache
 
 from pysetup.constants import (
     # code names
@@ -70,22 +70,26 @@ def installPackage(package: str):
 from marko.ext.gfm.elements import Table
 
 
+@lru_cache(maxsize=None)
 def _get_name_from_heading(heading: Heading) -> Optional[str]:
     last_child = heading.children[-1]
     if isinstance(last_child, CodeSpan):
         return last_child.children
     return None
 
 
+@lru_cache(maxsize=None)
 def _get_source_from_code_block(block: FencedCode) -> str:
     return block.children[0].children.strip()
 
 
+@lru_cache(maxsize=None)
 def _get_function_name_from_source(source: str) -> str:
     fn = ast.parse(source).body[0]
     return fn.name
 
 
+@lru_cache(maxsize=None)
 def _get_self_type_from_source(source: str) -> Optional[str]:
     fn = ast.parse(source).body[0]
     args = fn.args.args
@@ -98,6 +102,7 @@ def _get_self_type_from_source(source: str) -> Optional[str]:
     return args[0].annotation.id
 
 
+@lru_cache(maxsize=None)
 def _get_class_info_from_source(source: str) -> Tuple[str, Optional[str]]:
     class_def = ast.parse(source).body[0]
     base = class_def.bases[0]
@@ -113,12 +118,14 @@ def _get_class_info_from_source(source: str) -> Tuple[str, Optional[str]]:
     return class_def.name, parent_class
 
 
+@lru_cache(maxsize=None)
 def _is_constant_id(name: str) -> bool:
     if name[0] not in string.ascii_uppercase + '_':
         return False
     return all(map(lambda c: c in string.ascii_uppercase + '_' + string.digits, name[1:]))
 
 
+@lru_cache(maxsize=None)
 def _load_kzg_trusted_setups(preset_name):
     trusted_setups_file_path = str(Path(__file__).parent) + '/presets/' + preset_name + '/trusted_setups/trusted_setup_4096.json'
 
@@ -130,6 +137,7 @@ def _load_kzg_trusted_setups(preset_name):
 
     return trusted_setup_G1_monomial, trusted_setup_G1_lagrange, trusted_setup_G2_monomial
 
+@lru_cache(maxsize=None)
 def _load_curdleproofs_crs(preset_name):
     """
     NOTE: File generated from https://github.com/asn-d6/curdleproofs/blob/8e8bf6d4191fb6a844002f75666fb7009716319b/tests/crs.rs#L53-L67
@@ -153,6 +161,7 @@ def _load_curdleproofs_crs(preset_name):
 }
 
 
+@lru_cache(maxsize=None)
 def _get_eth2_spec_comment(child: LinkRefDef) -> Optional[str]:
     _, _, title = child._parse_info
     if not (title[0] == "(" and title[len(title)-1] == ")"):
@@ -163,6 +172,7 @@ def _get_eth2_spec_comment(child: LinkRefDef) -> Optional[str]:
     return title[len(ETH2_SPEC_COMMENT_PREFIX):].strip()
 
 
+@lru_cache(maxsize=None)
 def _parse_value(name: str, typed_value: str, type_hint: Optional[str] = None) -> VariableDefinition:
     comment = None
     if name in ("ROOT_OF_UNITY_EXTENDED", "ROOTS_OF_UNITY_EXTENDED", "ROOTS_OF_UNITY_REDUCED"):
@@ -185,6 +195,11 @@ def _update_constant_vars_with_kzg_setups(constant_vars, preset_name):
     constant_vars['KZG_SETUP_G2_MONOMIAL'] = VariableDefinition(constant_vars['KZG_SETUP_G2_MONOMIAL'].value, str(kzg_setups[2]), comment, None)
 
 
+@lru_cache(maxsize=None)
+def parse_markdown(content: str):
+    return gfm.parse(content)
+
+
 def get_spec(file_name: Path, preset: Dict[str, str], config: Dict[str, str], preset_name=str) -> SpecObject:
     functions: Dict[str, str] = {}
     protocols: Dict[str, ProtocolDefinition] = {}
@@ -198,7 +213,7 @@ def get_spec(file_name: Path, preset: Dict[str, str], config: Dict[str, str], pr
     custom_types: Dict[str, str] = {}
 
     with open(file_name) as source_file:
-        document = gfm.parse(source_file.read())
+        document = parse_markdown(source_file.read())
 
     current_name = None
     should_skip = False
@@ -326,6 +341,7 @@ def get_spec(file_name: Path, preset: Dict[str, str], config: Dict[str, str], pr
     )
 
 
+@lru_cache(maxsize=None)
 def load_preset(preset_files: Sequence[Path]) -> Dict[str, str]:
     """
     Loads the a directory of preset files, merges the result into one preset.
@@ -344,6 +360,7 @@ def load_preset(preset_files: Sequence[Path]) -> Dict[str, str]:
     return parse_config_vars(preset)
 
 
+@lru_cache(maxsize=None)
 def load_config(config_path: Path) -> Dict[str, str]:
     """
     Loads the given configuration file.
@@ -358,7 +375,7 @@ def build_spec(fork: str,
                source_files: Sequence[Path],
                preset_files: Sequence[Path],
                config_file: Path) -> str:
-    preset = load_preset(preset_files)
+    preset = load_preset(tuple(preset_files))
     config = load_config(config_file)
     all_specs = [get_spec(spec, preset, config, preset_name) for spec in source_files]