From 2e4233274f9b09a37085fd3802d08f00db152c1c Mon Sep 17 00:00:00 2001
From: Sebastien Bourdeauducq <sb@m-labs.hk>
Date: Mon, 28 Feb 2022 11:53:41 +0800
Subject: [PATCH] grabber: port to NAC3

---
 artiq/coredevice/grabber.py                 | 22 ++++++++++++---------
 artiq/examples/nac3devices/nac3devices.json |  4 ++++
 artiq/examples/nac3devices/nac3devices.py   |  3 +++
 3 files changed, 20 insertions(+), 9 deletions(-)

diff --git a/artiq/coredevice/grabber.py b/artiq/coredevice/grabber.py
index 6f29debe3..c71414a6d 100644
--- a/artiq/coredevice/grabber.py
+++ b/artiq/coredevice/grabber.py
@@ -1,19 +1,23 @@
 from numpy import int32, int64
 
 from artiq.language.core import *
-from artiq.language.types import *
 from artiq.coredevice.rtio import rtio_output, rtio_input_data
+from artiq.coredevice.core import Core
 
 
+@nac3
 class OutOfSyncException(Exception):
     """Raised when an incorrect number of ROI engine outputs has been
     retrieved from the RTIO input FIFO."""
     pass
 
 
+@nac3
 class Grabber:
     """Driver for the Grabber camera interface."""
-    kernel_invariants = {"core", "channel_base", "sentinel"}
+    core: KernelInvariant[Core]
+    channel_base: KernelInvariant[int32]
+    sentinel: KernelInvariant[int32]
 
     def __init__(self, dmgr, channel_base, res_width=12, count_shift=0,
                  core_device="core"):
@@ -26,7 +30,7 @@ class Grabber:
         self.sentinel = int32(int64(2**count_width))
 
     @kernel
-    def setup_roi(self, n, x0, y0, x1, y1):
+    def setup_roi(self, n: int32, x0: int32, y0: int32, x1: int32, y1: int32):
         """
         Defines the coordinates of a ROI.
 
@@ -50,7 +54,7 @@ class Grabber:
         delay_mu(c)
 
     @kernel
-    def gate_roi(self, mask):
+    def gate_roi(self, mask: int32):
         """
         Defines which ROI engines produce input events.
 
@@ -70,15 +74,15 @@ class Grabber:
         rtio_output((self.channel_base + 1) << 8, mask)
 
     @kernel
-    def gate_roi_pulse(self, mask, dt):
+    def gate_roi_pulse(self, mask: int32, dt: float):
         """Sets a temporary mask for the specified duration (in seconds), before
         disabling all ROI engines."""
         self.gate_roi(mask)
-        delay(dt)
+        self.core.delay(dt)
         self.gate_roi(0)
 
     @kernel
-    def input_mu(self, data):
+    def input_mu(self, data: list[int32]):
         """
         Retrieves the accumulated values for one frame from the ROI engines.
         Blocks until values are available.
@@ -96,10 +100,10 @@ class Grabber:
 
         sentinel = rtio_input_data(channel)
         if sentinel != self.sentinel:
-            raise OutOfSyncException
+            raise OutOfSyncException()
 
         for i in range(len(data)):
             roi_output = rtio_input_data(channel)
             if roi_output == self.sentinel:
-                raise OutOfSyncException
+                raise OutOfSyncException()
             data[i] = roi_output
diff --git a/artiq/examples/nac3devices/nac3devices.json b/artiq/examples/nac3devices/nac3devices.json
index a376ee3af..aae7184fe 100644
--- a/artiq/examples/nac3devices/nac3devices.json
+++ b/artiq/examples/nac3devices/nac3devices.json
@@ -32,6 +32,10 @@
             "bank_direction_low": "input",
             "bank_direction_high": "output",
             "edge_counter": true
+        },
+        {
+            "type": "grabber",
+            "ports": [6]
         }
     ]
 }
diff --git a/artiq/examples/nac3devices/nac3devices.py b/artiq/examples/nac3devices/nac3devices.py
index 051e56e0f..b979bd19c 100644
--- a/artiq/examples/nac3devices/nac3devices.py
+++ b/artiq/examples/nac3devices/nac3devices.py
@@ -8,6 +8,7 @@ from artiq.coredevice.urukul import CPLD as UrukulCPLD
 from artiq.coredevice.ad9912 import AD9912
 from artiq.coredevice.sampler import Sampler
 from artiq.coredevice.edge_counter import EdgeCounter
+from artiq.coredevice.grabber import Grabber
 
 
 @nac3
@@ -21,6 +22,7 @@ class NAC3Devices(EnvExperiment):
     urukul0_ch0: KernelInvariant[AD9912]
     sampler0: KernelInvariant[Sampler]
     ttl0_counter: KernelInvariant[EdgeCounter]
+    grabber0: KernelInvariant[Grabber]
 
     def build(self):
         self.setattr_device("core")
@@ -32,6 +34,7 @@ class NAC3Devices(EnvExperiment):
         self.setattr_device("urukul0_ch0")
         self.setattr_device("sampler0")
         self.setattr_device("ttl0_counter")
+        self.setattr_device("grabber0")
 
     @kernel
     def run(self):