From 284382b1f5b601e13394263cbe567ecbbc573f83 Mon Sep 17 00:00:00 2001
From: whitequark <whitequark@whitequark.org>
Date: Fri, 9 Jun 2017 07:10:30 +0000
Subject: [PATCH] compiler: add support for bytearray values in RPC (#714).

---
 artiq/compiler/transforms/llvm_ir_generator.py |  2 ++
 artiq/coredevice/comm_kernel.py                |  6 ++++++
 artiq/firmware/libproto/rpc_proto.rs           | 10 ++++++++--
 artiq/test/coredevice/test_embedding.py        |  3 +++
 4 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/artiq/compiler/transforms/llvm_ir_generator.py b/artiq/compiler/transforms/llvm_ir_generator.py
index cbcd45e85..8508d7043 100644
--- a/artiq/compiler/transforms/llvm_ir_generator.py
+++ b/artiq/compiler/transforms/llvm_ir_generator.py
@@ -1170,6 +1170,8 @@ class LLVMIRGenerator:
             return b"s"
         elif builtins.is_bytes(typ):
             return b"B"
+        elif builtins.is_bytearray(typ):
+            return b"A"
         elif builtins.is_list(typ):
             return b"l" + self._rpc_tag(builtins.get_iterable_elt(typ),
                                         error_handler)
diff --git a/artiq/coredevice/comm_kernel.py b/artiq/coredevice/comm_kernel.py
index d75c8e96a..66401f64e 100644
--- a/artiq/coredevice/comm_kernel.py
+++ b/artiq/coredevice/comm_kernel.py
@@ -371,6 +371,8 @@ class CommKernel:
             return self._read_string()
         elif tag == "B":
             return self._read_bytes()
+        elif tag == "A":
+            return self._read_bytes()
         elif tag == "l":
             length = self._read_int32()
             return [self._receive_rpc_value(embedding_map) for _ in range(length)]
@@ -467,6 +469,10 @@ class CommKernel:
             check(isinstance(value, bytes),
                   lambda: "bytes")
             self._write_bytes(value)
+        elif tag == "A":
+            check(isinstance(value, bytearray),
+                  lambda: "bytearray")
+            self._write_bytes(value)
         elif tag == "l":
             check(isinstance(value, list),
                   lambda: "list")
diff --git a/artiq/firmware/libproto/rpc_proto.rs b/artiq/firmware/libproto/rpc_proto.rs
index 51984b8f0..bd1e190f4 100644
--- a/artiq/firmware/libproto/rpc_proto.rs
+++ b/artiq/firmware/libproto/rpc_proto.rs
@@ -28,7 +28,7 @@ unsafe fn recv_value(reader: &mut Read, tag: Tag, data: &mut *mut (),
             consume_value!(u64, |ptr| {
                 *ptr = reader.read_u64()?; Ok(())
             }),
-        Tag::String | Tag::Bytes => {
+        Tag::String | Tag::Bytes | Tag::ByteArray => {
             consume_value!(CMutSlice<u8>, |ptr| {
                 let length = reader.read_u32()? as usize;
                 *ptr = CMutSlice::new(alloc(length)? as *mut u8, length);
@@ -108,7 +108,7 @@ unsafe fn send_value(writer: &mut Write, tag: Tag, data: &mut *const ()) -> io::
         Tag::String =>
             consume_value!(CSlice<u8>, |ptr|
                 writer.write_string(str::from_utf8((*ptr).as_ref()).unwrap())),
-        Tag::Bytes =>
+        Tag::Bytes | Tag::ByteArray =>
             consume_value!(CSlice<u8>, |ptr|
                 writer.write_bytes((*ptr).as_ref())),
         Tag::Tuple(it, arity) => {
@@ -206,6 +206,7 @@ mod tag {
         Float64,
         String,
         Bytes,
+        ByteArray,
         Tuple(TagIterator<'a>, u8),
         List(TagIterator<'a>),
         Array(TagIterator<'a>),
@@ -224,6 +225,7 @@ mod tag {
                 Tag::Float64 => b'f',
                 Tag::String => b's',
                 Tag::Bytes => b'B',
+                Tag::ByteArray => b'A',
                 Tag::Tuple(_, _) => b't',
                 Tag::List(_) => b'l',
                 Tag::Array(_) => b'a',
@@ -242,6 +244,7 @@ mod tag {
                 Tag::Float64 => 8,
                 Tag::String => 4,
                 Tag::Bytes => 4,
+                Tag::ByteArray => 4,
                 Tag::Tuple(it, arity) => {
                     let mut size = 0;
                     for _ in 0..arity {
@@ -287,6 +290,7 @@ mod tag {
                 b'f' => Tag::Float64,
                 b's' => Tag::String,
                 b'B' => Tag::Bytes,
+                b'A' => Tag::ByteArray,
                 b't' => {
                     let count = self.data[0];
                     self.data = &self.data[1..];
@@ -336,6 +340,8 @@ mod tag {
                         write!(f, "String")?,
                     Tag::Bytes =>
                         write!(f, "Bytes")?,
+                    Tag::ByteArray =>
+                        write!(f, "ByteArray")?,
                     Tag::Tuple(it, _) => {
                         write!(f, "Tuple(")?;
                         it.fmt(f)?;
diff --git a/artiq/test/coredevice/test_embedding.py b/artiq/test/coredevice/test_embedding.py
index 87f6dad77..bf692311b 100644
--- a/artiq/test/coredevice/test_embedding.py
+++ b/artiq/test/coredevice/test_embedding.py
@@ -41,6 +41,9 @@ class RoundtripTest(ExperimentCase):
     def test_bytes(self):
         self.assertRoundtrip(b"foo")
 
+    def test_bytearray(self):
+        self.assertRoundtrip(bytearray(b"foo"))
+
     def test_list(self):
         self.assertRoundtrip([10])