From bf84226c7d19b4bceecaf2842391b6dcef5828b6 Mon Sep 17 00:00:00 2001
From: David Nadlinger <code@klickverbot.at>
Date: Mon, 25 Jun 2018 10:52:25 +0100
Subject: [PATCH] language: Support appending to datasets

---
 RELEASE_NOTES.rst             |  2 ++
 artiq/language/environment.py | 12 ++++++++++++
 artiq/master/worker_db.py     | 16 ++++++++++------
 artiq/test/test_datasets.py   | 35 +++++++++++++++++++++++++++++++++++
 4 files changed, 59 insertions(+), 6 deletions(-)

diff --git a/RELEASE_NOTES.rst b/RELEASE_NOTES.rst
index a3beab361..453c4ed2b 100644
--- a/RELEASE_NOTES.rst
+++ b/RELEASE_NOTES.rst
@@ -20,6 +20,8 @@ ARTIQ-5
   edge timestamps are not required. See :mod:`artiq.coredevice.edge_counter` for
   the core device driver and :mod:`artiq.gateware.rtio.phy.edge_counter`/
   :meth:`artiq.gateware.eem.DIO.add_std` for the gateware components.
+* List datasets can now be efficiently appended to from experiments using
+  :meth:`artiq.language.environment.HasEnvironment.append_to_dataset`.
 * The controller manager now ignores device database entries without the
   ``"command"`` key set to facilitate sharing of devices between multiple
   masters.
diff --git a/artiq/language/environment.py b/artiq/language/environment.py
index 5eb4e8d81..0e1b2b918 100644
--- a/artiq/language/environment.py
+++ b/artiq/language/environment.py
@@ -321,6 +321,18 @@ class HasEnvironment:
         as ``slice(*sub_tuple)`` (multi-dimensional slicing)."""
         self.__dataset_mgr.mutate(key, index, value)
 
+    @rpc(flags={"async"})
+    def append_to_dataset(self, key, value):
+        """Append a value to a dataset.
+
+        The target dataset must be a list (i.e. support ``append()``), and must
+        have previously been set from this experiment.
+
+        The broadcast/persist/archive mode of the given key remains unchanged
+        from when the dataset was last set. Appended values are transmitted
+        efficiently as incremental modifications in broadcast mode."""
+        self.__dataset_mgr.append_to(key, value)
+
     def get_dataset(self, key, default=NoDefault, archive=True):
         """Returns the contents of a dataset.
 
diff --git a/artiq/master/worker_db.py b/artiq/master/worker_db.py
index a1555037f..d2dd628f5 100644
--- a/artiq/master/worker_db.py
+++ b/artiq/master/worker_db.py
@@ -136,17 +136,18 @@ class DatasetManager:
         elif key in self.local:
             del self.local[key]
 
-    def mutate(self, key, index, value):
-        target = None
-        if key in self.local:
-            target = self.local[key]
+    def _get_mutation_target(self, key):
+        target = self.local.get(key, None)
         if key in self._broadcaster.raw_view:
             if target is not None:
                 assert target is self._broadcaster.raw_view[key][1]
-            target = self._broadcaster[key][1]
+            return self._broadcaster[key][1]
         if target is None:
-            raise KeyError("Cannot mutate non-existing dataset")
+            raise KeyError("Cannot mutate nonexistent dataset '{}'".format(key))
+        return target
 
+    def mutate(self, key, index, value):
+        target = self._get_mutation_target(key)
         if isinstance(index, tuple):
             if isinstance(index[0], tuple):
                 index = tuple(slice(*e) for e in index)
@@ -154,6 +155,9 @@ class DatasetManager:
                 index = slice(*index)
         setitem(target, index, value)
 
+    def append_to(self, key, value):
+        self._get_mutation_target(key).append(value)
+
     def get(self, key, archive=False):
         if key in self.local:
             return self.local[key]
diff --git a/artiq/test/test_datasets.py b/artiq/test/test_datasets.py
index 710e34116..db35d7f34 100644
--- a/artiq/test/test_datasets.py
+++ b/artiq/test/test_datasets.py
@@ -32,6 +32,9 @@ class TestExperiment(EnvExperiment):
     def set(self, key, value, **kwargs):
         self.set_dataset(key, value, **kwargs)
 
+    def append(self, key, value):
+        self.append_to_dataset(key, value)
+
 
 KEY = "foo"
 
@@ -67,3 +70,35 @@ class ExperimentDatasetCase(unittest.TestCase):
         self.assertEqual(self.exp.get(KEY), 1)
         with self.assertRaises(KeyError):
             self.dataset_db.get(KEY)
+
+    def test_append_local(self):
+        self.exp.set(KEY, [])
+        self.exp.append(KEY, 0)
+        self.assertEqual(self.exp.get(KEY), [0])
+        self.exp.append(KEY, 1)
+        self.assertEqual(self.exp.get(KEY), [0, 1])
+
+    def test_append_broadcast(self):
+        self.exp.set(KEY, [], broadcast=True)
+        self.exp.append(KEY, 0)
+        self.assertEqual(self.dataset_db.data[KEY][1], [0])
+        self.exp.append(KEY, 1)
+        self.assertEqual(self.dataset_db.data[KEY][1], [0, 1])
+
+    def test_append_array(self):
+        for broadcast in (True, False):
+            self.exp.set(KEY, [], broadcast=broadcast)
+            self.exp.append(KEY, [])
+            self.exp.append(KEY, [])
+            self.assertEqual(self.exp.get(KEY), [[], []])
+
+    def test_append_scalar_fails(self):
+        for broadcast in (True, False):
+            with self.assertRaises(AttributeError):
+                self.exp.set(KEY, 0, broadcast=broadcast)
+                self.exp.append(KEY, 1)
+
+    def test_append_nonexistent_fails(self):
+        with self.assertRaises(KeyError):
+            self.exp.append(KEY, 0)
+