From 7385b911130f6684e4fdd83e8009a767048543ec Mon Sep 17 00:00:00 2001
From: Sebastien Bourdeauducq <sb@m-labs.hk>
Date: Sat, 6 Nov 2021 18:41:59 +0800
Subject: [PATCH] nac3artiq: support kernel entry short form from original
 ARTIQ

---
 nac3artiq/demo/demo.py      |  5 +----
 nac3artiq/demo/min_artiq.py | 22 +++++++++++++++-------
 2 files changed, 16 insertions(+), 11 deletions(-)

diff --git a/nac3artiq/demo/demo.py b/nac3artiq/demo/demo.py
index 2a4b84907..75c328e95 100644
--- a/nac3artiq/demo/demo.py
+++ b/nac3artiq/demo/demo.py
@@ -12,7 +12,7 @@ class Demo:
         self.led1 = TTLOut(self.core, 19)
 
     @kernel
-    def run_k(self):
+    def run(self):
         self.core.reset()
         while True:
             with parallel:
@@ -20,9 +20,6 @@ class Demo:
                 self.led1.pulse(100.*ms)
             self.core.delay(100.*ms)
 
-    def run(self):
-        self.core.run(self.run_k)
-
 
 if __name__ == "__main__":
     Demo().run()
diff --git a/nac3artiq/demo/min_artiq.py b/nac3artiq/demo/min_artiq.py
index 06c3640d6..c7913b7a2 100644
--- a/nac3artiq/demo/min_artiq.py
+++ b/nac3artiq/demo/min_artiq.py
@@ -1,6 +1,7 @@
-from inspect import isclass, getmodule
+from inspect import isclass, getfullargspec, getmodule
 from functools import wraps
 import sys
+from types import SimpleNamespace
 from numpy import int32, int64
 
 import nac3artiq
@@ -37,13 +38,20 @@ def extern(function):
     return function
 
 
-def kernel(class_or_function):
+def kernel(function_or_method):
     """Decorates a function or method to be executed on the core device."""
-    register_module_of(class_or_function)
-    @wraps(class_or_function)
-    def device_only(*args, **kwargs):
-        raise RuntimeError("Kernels must not be called directly, use core.run(kernel_function) instead")
-    return device_only
+    register_module_of(function_or_method)
+    argspec = getfullargspec(function_or_method)
+    if argspec.args and argspec.args[0] == "self":
+        @wraps(function_or_method)
+        def run_on_core(self, *args, **kwargs):
+            fake_method = SimpleNamespace(__self__=self, __name__=function_or_method.__name__)
+            self.core.run(fake_method, *args, **kwargs)
+    else:
+        @wraps(function_or_method)
+        def run_on_core(*args, **kwargs):
+            raise RuntimeError("Kernel functions need explicit core.run()")
+    return run_on_core
 
 
 def portable(function):