From 84f472501598120aeece42d9a696f739946dec8c Mon Sep 17 00:00:00 2001
From: Sebastien Bourdeauducq <sb@m-labs.hk>
Date: Sat, 6 Aug 2016 12:01:49 +0800
Subject: [PATCH] cache source on import of modules that may contain kernels.
 Closes #416

---
 artiq/compiler/import_cache.py | 55 ++++++++++++++++++++++++++++++++++
 artiq/frontend/artiq_run.py    |  2 ++
 artiq/language/core.py         |  6 ++++
 artiq/master/worker_impl.py    |  3 ++
 4 files changed, 66 insertions(+)
 create mode 100644 artiq/compiler/import_cache.py

diff --git a/artiq/compiler/import_cache.py b/artiq/compiler/import_cache.py
new file mode 100644
index 000000000..6a092e447
--- /dev/null
+++ b/artiq/compiler/import_cache.py
@@ -0,0 +1,55 @@
+import sys
+import builtins
+import linecache
+import tokenize
+import logging
+import importlib.machinery as im
+
+from artiq.experiment import kernel, portable
+
+
+__all__ = ["install_hook"]
+
+
+logger = logging.getLogger(__name__)
+
+
+cache = dict()
+im_exec_module = None
+linecache_getlines = None
+
+
+def hook_exec_module(self, module):
+    im_exec_module(self, module)
+    if (hasattr(module, "__file__")
+            # Heuristic to determine if the module may contain ARTIQ kernels.
+            # This breaks if kernel is not imported the usual way.
+            and ((getattr(module, "kernel", None) is kernel)
+                 or (getattr(module, "portable", None) is portable))):
+        fn = module.__file__
+        try:
+            with tokenize.open(fn) as fp:
+                cache[fn] = fp.readlines()
+        except:
+            logger.warning("failed to add '%s' to cache", fn, exc_info=True)
+        else:
+            logger.debug("added '%s' to cache", fn)
+
+
+def hook_getlines(filename, module_globals=None):
+    if filename in cache:
+        return cache[filename]
+    else:
+        return linecache_getlines(filename, module_globals)
+
+
+def install_hook():
+    global im_exec_module, linecache_getlines
+
+    im_exec_module = im.SourceFileLoader.exec_module
+    im.SourceFileLoader.exec_module = hook_exec_module
+
+    linecache_getlines = linecache.getlines
+    linecache.getlines = hook_getlines
+
+    logger.debug("hook installed")
diff --git a/artiq/frontend/artiq_run.py b/artiq/frontend/artiq_run.py
index a1a0323cd..09b46ea07 100755
--- a/artiq/frontend/artiq_run.py
+++ b/artiq/frontend/artiq_run.py
@@ -19,6 +19,7 @@ from artiq.master.worker_db import DeviceManager, DatasetManager
 from artiq.coredevice.core import CompileError, host_only
 from artiq.compiler.embedding import EmbeddingMap
 from artiq.compiler.targets import OR1KTarget
+from artiq.compiler import import_cache
 from artiq.tools import *
 
 
@@ -158,6 +159,7 @@ def _build_experiment(device_mgr, dataset_mgr, args):
         elif is_bc:
             return LLVMBitcodeRunner(device_mgr, dataset_mgr, file=args.file)
         else:
+            import_cache.install_hook()
             module = file_import(args.file, prefix="artiq_run_")
         file = args.file
     else:
diff --git a/artiq/language/core.py b/artiq/language/core.py
index 85660ab55..b08da2147 100644
--- a/artiq/language/core.py
+++ b/artiq/language/core.py
@@ -43,6 +43,9 @@ def kernel(arg=None, flags={}):
 
     The decorator takes an optional parameter that defaults to ``core`` and
     specifies the name of the attribute to use as core device driver.
+
+    This decorator must be present in the global namespace of all modules using
+    it for the import cache to work properly.
     """
     if isinstance(arg, str):
         def inner_decorator(function):
@@ -70,6 +73,9 @@ def portable(arg=None, flags={}):
     host will be executed on the host (no compilation and execution on the
     core device). A decorated function called from a kernel will be executed
     on the core device (no RPC).
+
+    This decorator must be present in the global namespace of all modules using
+    it for the import cache to work properly.
     """
     if arg is None:
         def inner_decorator(function):
diff --git a/artiq/master/worker_impl.py b/artiq/master/worker_impl.py
index 6c23c29fc..c52f09226 100644
--- a/artiq/master/worker_impl.py
+++ b/artiq/master/worker_impl.py
@@ -16,6 +16,7 @@ from artiq.language.environment import (is_experiment, TraceArgumentManager,
                                         ProcessArgumentManager)
 from artiq.language.core import set_watchdog_factory, TerminationRequested
 from artiq.language.types import TBool
+from artiq.compiler import import_cache
 from artiq.coredevice.core import CompileError, host_only, _render_diagnostic
 from artiq import __version__ as artiq_version
 
@@ -191,6 +192,8 @@ def main():
                                virtual_devices={"scheduler": Scheduler()})
     dataset_mgr = DatasetManager(ParentDatasetDB)
 
+    import_cache.install_hook()
+
     try:
         while True:
             obj = get_object()