From 34cf303e6c1fa153f3b712dd5f42229578263abb Mon Sep 17 00:00:00 2001
From: pca006132 <jl@m-labs.hk>
Date: Sat, 6 Nov 2021 22:50:28 +0800
Subject: [PATCH] nac3artiq: modified demo to use KernelInvariants

---
 nac3artiq/demo/demo.py      |  6 +++---
 nac3artiq/demo/min_artiq.py | 16 +++++++++-------
 2 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/nac3artiq/demo/demo.py b/nac3artiq/demo/demo.py
index 75c328e9..a9abc18b 100644
--- a/nac3artiq/demo/demo.py
+++ b/nac3artiq/demo/demo.py
@@ -2,9 +2,9 @@ from min_artiq import *
 
 @nac3
 class Demo:
-    core: Core
-    led0: TTLOut
-    led1: TTLOut
+    core: KernelInvariant[Core]
+    led0: KernelInvariant[TTLOut]
+    led1: KernelInvariant[TTLOut]
 
     def __init__(self):
         self.core = Core()
diff --git a/nac3artiq/demo/min_artiq.py b/nac3artiq/demo/min_artiq.py
index c7913b7a..800dc65d 100644
--- a/nac3artiq/demo/min_artiq.py
+++ b/nac3artiq/demo/min_artiq.py
@@ -3,6 +3,7 @@ from functools import wraps
 import sys
 from types import SimpleNamespace
 from numpy import int32, int64
+from typing import Generic, TypeVar
 
 import nac3artiq
 
@@ -18,8 +19,10 @@ compiler = nac3artiq.NAC3(core_arguments["target"])
 allow_module_registration = True
 registered_ids = set()
 
-def KernelInvariant(t):
-    return t
+
+T = TypeVar('T')
+class KernelInvariant(Generic[T]):
+    pass
 
 
 def register_module_of(obj):
@@ -100,7 +103,7 @@ def rtio_input_data(channel: int32) -> int32:
 
 @nac3
 class Core:
-    ref_period: float
+    ref_period: KernelInvariant[float]
 
     def __init__(self):
         self.ref_period = core_arguments["ref_period"]
@@ -146,11 +149,10 @@ class Core:
 
 @nac3
 class TTLOut:
-    core: Core
-    channel: int32
-    target_o: int32
+    core: KernelInvariant[Core]
+    channel: KernelInvariant[int32]
+    target_o: KernelInvariant[int32]
 
-    @portable
     def __init__(self, core: Core, channel: int32):
         self.core = core
         self.channel = channel