Implement bytecode caching in BasilispImporter (#244)

* Implement bytecode caching in BasilispImporter

* Make sure to bootstrap the cached module first

* Time import time

* Remove unused single quote

* Fix MyPy and Lint errors

* Factor out into multiple methods to suppress linting errors

Author: chrisrink10

src/basilisp/compiler.py

@@ -1830,11 +1830,15 @@ def to_py_str(t: ast.AST) -> str:
     return codegen.to_source(t)
 
 
-def compile_and_exec_form(form: LispForm,
+BytecodeCollector = Optional[Callable[[types.CodeType], None]]
+
+
+def compile_and_exec_form(form: LispForm,  # pylint: disable= too-many-arguments
                           ctx: CompilerContext,
                           module: types.ModuleType,
                           source_filename: str = '<REPL Input>',
-                          wrapped_fn_name: str = _DEFAULT_FN) -> Any:
+                          wrapped_fn_name: str = _DEFAULT_FN,
+                          collect_bytecode: Optional[BytecodeCollector] = None) -> Any:
     """Compile and execute the given form. This function will be most useful
     for the REPL and testing purposes. Returns the result of the executed expression.
 
@@ -1867,13 +1871,16 @@ def compile_and_exec_form(form: LispForm,
         runtime.add_generated_python(to_py_str(ast_module))
 
     bytecode = compile(ast_module, source_filename, 'exec')
+    if collect_bytecode:
+        collect_bytecode(bytecode)
     exec(bytecode, module.__dict__)
     return getattr(module, final_wrapped_name)()
 
 
 def _incremental_compile_module(nodes: MixedNodeStream,
                                 mod: types.ModuleType,
-                                source_filename: str) -> None:
+                                source_filename: str,
+                                collect_bytecode: Optional[BytecodeCollector] = None) -> None:
     """Incrementally compile a stream of AST nodes in module mod.
 
     The source_filename will be passed to Python's native compile.
@@ -1891,24 +1898,31 @@ def _incremental_compile_module(nodes: MixedNodeStream,
         runtime.add_generated_python(to_py_str(module))
 
     bytecode = compile(module, source_filename, 'exec')
+    if collect_bytecode:
+        collect_bytecode(bytecode)
     exec(bytecode, mod.__dict__)
 
 
-def _bootstrap_module(ctx: CompilerContext, mod: types.ModuleType, source_filename: str) -> None:
+def _bootstrap_module(ctx: CompilerContext,
+                      mod: types.ModuleType,
+                      source_filename: str,
+                      collect_bytecode: Optional[BytecodeCollector] = None) -> None:
     """Bootstrap a new module with imports and other boilerplate."""
     preamble: List[ast.AST] = []
     preamble.extend(_module_imports(ctx))
     preamble.append(_from_module_import())
     preamble.append(_ns_var())
 
-    _incremental_compile_module(preamble, mod, source_filename=source_filename)
+    _incremental_compile_module(
+        preamble, mod, source_filename=source_filename, collect_bytecode=collect_bytecode)
     mod.__basilisp_bootstrapped__ = True  # type: ignore
 
 
 def compile_module(forms: Iterable[LispForm],
                    ctx: CompilerContext,
                    module: types.ModuleType,
-                   source_filename: str) -> None:
+                   source_filename: str,
+                   collect_bytecode: Optional[BytecodeCollector] = None) -> None:
     """Compile an entire Basilisp module into Python bytecode which can be
     executed as a Python module.
 
@@ -1920,7 +1934,23 @@ def compile_module(forms: Iterable[LispForm],
 
     for form in forms:
         nodes = [node for node in _to_ast(ctx, form)]
-        _incremental_compile_module(nodes, module, source_filename=source_filename)
+        _incremental_compile_module(
+            nodes, module, source_filename=source_filename, collect_bytecode=collect_bytecode)
+
+
+def compile_bytecode(code: List[types.CodeType],
+                     ctx: CompilerContext,
+                     module: types.ModuleType,
+                     source_filename: str) -> None:
+    """Compile cached bytecode into the given module.
+
+    The Basilisp import hook attempts to cache bytecode while compiling Basilisp
+    namespaces. When the cached bytecode is reloaded from disk, it needs to be
+    compiled within a bootstrapped module. This function bootstraps the module
+    and then proceeds to compile a collection of bytecodes into the module."""
+    _bootstrap_module(ctx, module, source_filename)
+    for bytecode in code:
+        exec(bytecode, module.__dict__)
 
 
 lrepr = basilisp.lang.util.lrepr

src/basilisp/importer.py

@@ -1,20 +1,89 @@
 import importlib.machinery
 import importlib.util
 import logging
+import marshal
+import os
 import os.path
 import sys
 import types
 from importlib.abc import MetaPathFinder, SourceLoader
-from typing import Optional
+from typing import Optional, List, Dict
 
 import basilisp.compiler as compiler
 import basilisp.lang.runtime as runtime
 import basilisp.reader as reader
 from basilisp.lang.util import demunge
+from basilisp.util import timed
+
+MAGIC_NUMBER = (1149).to_bytes(2, 'little') + b'\r\n'
 
 logger = logging.getLogger(__name__)
 
 
+def _r_long(int_bytes: bytes) -> int:
+    """Convert 4 bytes in little-endian to an integer."""
+    return int.from_bytes(int_bytes, 'little')
+
+
+def _w_long(x: int) -> bytes:
+    """Convert a 32-bit integer to little-endian."""
+    return (int(x) & 0xFFFFFFFF).to_bytes(4, 'little')
+
+
+def _basilisp_bytecode(mtime: int,
+                       source_size: int,
+                       code: List[types.CodeType]) -> bytes:
+    """Return the bytes for a Basilisp bytecode cache file."""
+    data = bytearray(MAGIC_NUMBER)
+    data.extend(_w_long(mtime))
+    data.extend(_w_long(source_size))
+    data.extend(marshal.dumps(code))  # type: ignore
+    return data
+
+
+def _get_basilisp_bytecode(fullname: str,
+                           mtime: int,
+                           source_size: int,
+                           cache_data: bytes) -> List[types.CodeType]:
+    """Unmarshal the bytes from a Basilisp bytecode cache file, validating the
+    file header prior to returning. If the file header does not match, throw
+    an exception."""
+    exc_details = {'name': fullname}
+    magic = cache_data[:4]
+    raw_timestamp = cache_data[4:8]
+    raw_size = cache_data[8:12]
+    if magic != MAGIC_NUMBER:
+        message = f"Incorrect magic number ({magic}) in {fullname}; expected {MAGIC_NUMBER}"
+        logger.debug(message)
+        raise ImportError(message, **exc_details)  # type: ignore
+    elif len(raw_timestamp) != 4:
+        message = f"Reached EOF while reading timestamp in {fullname}"
+        logger.debug(message)
+        raise EOFError(message)
+    elif _r_long(raw_timestamp) != mtime:
+        message = f"Non-matching timestamp ({_r_long(raw_timestamp)}) in {fullname} bytecode cache; expected {mtime}"
+        logger.debug(message)
+        raise ImportError(message, **exc_details)  # type: ignore
+    elif len(raw_size) != 4:
+        message = f"Reached EOF while reading size of source in {fullname}"
+        logger.debug(message)
+        raise EOFError(message)
+    elif _r_long(raw_size) != source_size:
+        message = f"Non-matching filesize ({_r_long(raw_size)}) in {fullname} bytecode cache; expected {source_size}"
+        logger.debug(message)
+        raise ImportError(message, **exc_details)  # type: ignore
+
+    return marshal.loads(cache_data[12:])  # type: ignore
+
+
+def _cache_from_source(path: str) -> str:
+    """Return the path to the cached file for the given path. The original path
+    does not have to exist."""
+    cache_path, cache_file = os.path.split(importlib.util.cache_from_source(path))
+    filename, _ = os.path.splitext(cache_file)
+    return os.path.join(cache_path, filename + '.lpyc')
+
+
 class BasilispImporter(MetaPathFinder, SourceLoader):
     """Python import hook to allow directly loading Basilisp code within
     Python."""
@@ -41,7 +110,13 @@ def find_spec(self,
                          f"{os.path.join(entry, *module_name)}.lpy"]
             for filename in filenames:
                 if os.path.exists(filename):
-                    state = {'fullname': fullname, "filename": filename, 'path': entry, 'target': target}
+                    state = {
+                        'fullname': fullname,
+                        'filename': filename,
+                        'path': entry,
+                        'target': target,
+                        'cache_filename': _cache_from_source(filename)
+                    }
                     logger.debug(f"Found potential Basilisp module '{fullname}' in file '{filename}'")
                     return importlib.machinery.ModuleSpec(fullname, self, origin=filename, loader_state=state)
         return None
@@ -50,10 +125,22 @@ def invalidate_caches(self):
         super().invalidate_caches()
         self._cache = {}
 
-    def get_data(self, path) -> bytes:
+    def _cache_bytecode(self, source_path, cache_path, data):  # pylint: disable=unused-argument
+        self.set_data(cache_path, data)
+
+    def path_stats(self, path):
+        stat = os.stat(path)
+        return {'mtime': int(stat.st_mtime), 'size': stat.st_size}
+
+    def get_data(self, path):
         with open(path, mode='r+b') as f:
             return f.read()
 
+    def set_data(self, path, data):
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+        with open(path, mode='w+b') as f:
+            f.write(data)
+
     def get_filename(self, fullname: str) -> str:
         try:
             cached = self._cache[fullname]
@@ -72,6 +159,50 @@ def create_module(self, spec: importlib.machinery.ModuleSpec):
         self._cache[spec.name] = {"spec": spec}
         return mod
 
+    def _exec_cached_module(self,
+                            fullname: str,
+                            loader_state: Dict[str, str],
+                            path_stats: Dict[str, int],
+                            module: types.ModuleType):
+        """Load and execute a cached Basilisp module."""
+        filename = loader_state["filename"]
+        cache_filename = loader_state["cache_filename"]
+
+        with timed(lambda duration: logger.debug(
+                f"Loaded cached Basilisp module '{fullname}' in {duration / 1000000}ms")):
+            logger.debug(f"Checking for cached Basilisp module '{fullname}''")
+            cache_data = self.get_data(cache_filename)
+            cached_code = _get_basilisp_bytecode(fullname, path_stats['mtime'], path_stats['size'], cache_data)
+            compiler.compile_bytecode(cached_code, compiler.CompilerContext(), module, filename)
+
+    def _exec_module(self,
+                     fullname: str,
+                     loader_state: Dict[str, str],
+                     path_stats: Dict[str, int],
+                     module: types.ModuleType):
+        """Load and execute a non-cached Basilisp module."""
+        filename = loader_state["filename"]
+        cache_filename = loader_state["cache_filename"]
+
+        with timed(lambda duration: logger.debug(
+                f"Loaded Basilisp module '{fullname}' in {duration / 1000000}ms")):
+            # During compilation, bytecode objects are added to the list via the closure
+            # add_bytecode below, which is passed to the compiler. The collected bytecodes
+            # will be used to generate an .lpyc file for caching the compiled file.
+            all_bytecode = []
+
+            def add_bytecode(bytecode: types.CodeType):
+                all_bytecode.append(bytecode)
+
+            logger.debug(f"Reading and compiling Basilisp module '{fullname}'")
+            forms = reader.read_file(filename, resolver=runtime.resolve_alias)
+            compiler.compile_module(  # pylint: disable=unexpected-keyword-arg
+                forms, compiler.CompilerContext(), module, filename, collect_bytecode=add_bytecode)
+
+        # Cache the bytecode that was collected through the compilation run.
+        cache_file_bytes = _basilisp_bytecode(path_stats['mtime'], path_stats['size'], all_bytecode)
+        self._cache_bytecode(filename, cache_filename, cache_file_bytes)
+
     def exec_module(self, module):
         """Compile the Basilisp module into Python code.
 
@@ -84,6 +215,7 @@ def exec_module(self, module):
         cached["module"] = module
         spec = cached["spec"]
         filename = spec.loader_state["filename"]
+        path_stats = self.path_stats(filename)
 
         # During the bootstrapping process, the 'basilisp.core namespace is created with
         # a blank module. If we do not replace the module here with the module we are
@@ -93,10 +225,13 @@ def exec_module(self, module):
         ns: runtime.Namespace = runtime.set_current_ns(ns_name).value
         ns.module = module
 
-        logger.debug(f"Reading and compiling Basilisp module '{fullname}''")
-        forms = reader.read_file(filename, resolver=runtime.resolve_alias)
-        compiler.compile_module(forms, compiler.CompilerContext(), module, filename)
-        logger.debug(f"Loaded Basilisp module '{fullname}''")
+        # Check if a valid, cached version of this Basilisp namespace exists and, if so,
+        # load it and bypass the expensive compilation process below.
+        try:
+            self._exec_cached_module(fullname, spec.loader_state, path_stats, module)
+        except (EOFError, ImportError, IOError, OSError) as e:
+            logger.debug(f"Failed to load cached Basilisp module: {e}")
+            self._exec_module(fullname, spec.loader_state, path_stats, module)
 
         # Because we want to (by default) add 'basilisp.core into every namespace by default,
         # we want to make sure we don't try to add 'basilisp.core into itself, causing a

src/basilisp/util.py

@@ -1,6 +1,8 @@
+import contextlib
 import functools
 import inspect
 import os.path
+import time
 from typing import Optional, Callable, TypeVar, Generic
 
 from functional import seq
@@ -34,6 +36,18 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
+@contextlib.contextmanager
+def timed(f: Optional[Callable[[int], None]] = None):
+    """Time the execution of code in the with-block, calling the function
+    f (if it is given) with the resulting time in nanoseconds."""
+    start = time.perf_counter()
+    yield
+    end = time.perf_counter()
+    if f:
+        ns = int((end - start) * 1000000000)
+        f(ns)
+
+
 T = TypeVar('T')
 U = TypeVar('U')