From 80a366917ef4b9f9610156c1ab7b9d25398249e4 Mon Sep 17 00:00:00 2001
From: Simon Renblad <srenblad@m-labs.hk>
Date: Mon, 6 May 2024 17:30:14 +0800
Subject: [PATCH] working deferred args system

---
 artiq/language/environment.py | 38 ++++++++++++++++++++++++++++++++++-
 1 file changed, 37 insertions(+), 1 deletion(-)

diff --git a/artiq/language/environment.py b/artiq/language/environment.py
index a80701c78..d2e5c6592 100644
--- a/artiq/language/environment.py
+++ b/artiq/language/environment.py
@@ -228,20 +228,44 @@ class TraceArgumentManager:
     def get_interactive(self, interactive_arglist, title):
         raise NotImplementedError
 
+    def defer_processing(self, key):
+        pass
+
+    def get_deferred_arguments(self):
+        pass
+
 
 class ProcessArgumentManager:
     def __init__(self, unprocessed_arguments):
         self.unprocessed_arguments = unprocessed_arguments
         self._processed_arguments = set()
+        self._deferred_arguments = dict()
 
     def get(self, key, processor, group, tooltip):
-        if key in self.unprocessed_arguments:
+        unprocessed = key in self.unprocessed_arguments
+        if unprocessed and key in self._deferred_arguments:
+            self._deferred_arguments[key] = processor
+            r = processor.default()
+        elif unprocessed:
             r = processor.process(self.unprocessed_arguments[key])
             self._processed_arguments.add(key)
         else:
             r = processor.default()
         return r
 
+    def defer_processing(self, key):
+        self._deferred_arguments[key] = None
+
+    def get_deferred_arguments(self):
+        deferred = dict()
+        for key, proc in self._deferred_arguments.items():
+            if proc is None:
+                raise AttributeError("No processor provided for deferred argument: {}".format(key))
+            deferred[key] = proc.process(self.unprocessed_arguments[key])
+            self._processed_arguments.add(key)
+        self._deferred_arguments.clear()
+        return deferred
+
     def check_unprocessed_arguments(self):
         unprocessed = set(self.unprocessed_arguments.keys()) -\
                       self._processed_arguments
@@ -342,6 +366,18 @@ class HasEnvironment:
         kernel_invariants = getattr(self, "kernel_invariants", set())
         self.kernel_invariants = kernel_invariants | {key}
 
+    def defer_argument_processing(self, key):
+        self.__argument_mgr.defer_processing(key)
+
+    def set_deferred_arguments(self):
+        deferred = self.__argument_mgr.get_deferred_arguments()
+        if deferred is None:
+            return
+        kernel_invariants = getattr(self, "kernel_invariants", set())
+        for k, v in deferred.items():
+            setattr(self, k, v)
+            self.kernel_invariants = kernel_invariants | {k}
+
     @contextmanager
     def interactive(self, title=""):
         """Request arguments from the user interactively.