From 5c21649d10b5a9c9de11c398a5b3bc2af3aa8d16 Mon Sep 17 00:00:00 2001
From: occheung <dc@m-labs.hk>
Date: Thu, 19 Sep 2024 09:46:17 +0800
Subject: [PATCH] frontend: make coremgmt flash zynq-compatible

---
 artiq/coredevice/comm_mgmt.py    |  3 +-
 artiq/frontend/artiq_coremgmt.py | 68 ++++++++++++++++++--------------
 2 files changed, 40 insertions(+), 31 deletions(-)

diff --git a/artiq/coredevice/comm_mgmt.py b/artiq/coredevice/comm_mgmt.py
index b0ab94509..c9f466b59 100644
--- a/artiq/coredevice/comm_mgmt.py
+++ b/artiq/coredevice/comm_mgmt.py
@@ -208,7 +208,8 @@ class CommMgmt:
             for filename in bin_paths:
                 with open(filename, "rb") as fi:
                     bin_ = fi.read()
-                    image_buf.write(struct.pack(self.endian + "I", len(bin_)))
+                    if (len(bin_paths) > 1):
+                        image_buf.write(struct.pack(self.endian + "I", len(bin_)))
                     image_buf.write(bin_)
 
             crc = binascii.crc32(image_buf.getvalue())
diff --git a/artiq/frontend/artiq_coremgmt.py b/artiq/frontend/artiq_coremgmt.py
index b51861a95..e9cada16f 100755
--- a/artiq/frontend/artiq_coremgmt.py
+++ b/artiq/frontend/artiq_coremgmt.py
@@ -94,6 +94,9 @@ def get_argparser():
     t_flash = tools.add_parser("flash",
                                help="flash the running system")
 
+    p_zynq = t_flash.add_argument("-z", "--zynq", default=False,
+                                  help="target zynq device", action="store_true")
+
     p_directory = t_flash.add_argument("directory", metavar="DIRECTORY", type=str,
                                        help="directory that contains the binaries")
 
@@ -160,40 +163,45 @@ def main():
             mgmt.config_erase()
     
     if args.tool == "flash":
-        def artifact_path(this_binary_dir, *path_filename):
-            if args.srcbuild:
-                # source tree - use path elements to locate file
-                return os.path.join(this_binary_dir, *path_filename)
-            else:
-                # flat tree - all files in the same directory, discard path elements
-                *_, filename = path_filename
-                return os.path.join(this_binary_dir, filename)
+        if args.zynq:
+            boot = os.path.join(args.directory, "boot.bin")
+            bins = [ boot ]
+        else:
+            def artifact_path(this_binary_dir, *path_filename):
+                if args.srcbuild:
+                    # source tree - use path elements to locate file
+                    return os.path.join(this_binary_dir, *path_filename)
+                else:
+                    # flat tree - all files in the same directory, discard path elements
+                    *_, filename = path_filename
+                    return os.path.join(this_binary_dir, filename)
 
-        def convert_gateware(bit_filename):
-            bin_handle, bin_filename = tempfile.mkstemp(
-                prefix="artiq_", suffix="_" + os.path.basename(bit_filename))
-            with open(bit_filename, "rb") as bit_file, open(bin_handle, "wb") as bin_file:
-                bit2bin(bit_file, bin_file)
-            atexit.register(lambda: os.unlink(bin_filename))
-            return bin_filename
+            def convert_gateware(bit_filename):
+                bin_handle, bin_filename = tempfile.mkstemp(
+                    prefix="artiq_", suffix="_" + os.path.basename(bit_filename))
+                with open(bit_filename, "rb") as bit_file, open(bin_handle, "wb") as bin_file:
+                    bit2bin(bit_file, bin_file)
+                atexit.register(lambda: os.unlink(bin_filename))
+                return bin_filename
 
-        gateware = convert_gateware(
-            artifact_path(args.directory, "gateware", "top.bit"))
-        bootloader = artifact_path(args.directory, "software", "bootloader", "bootloader.bin")
+            gateware = convert_gateware(
+                artifact_path(args.directory, "gateware", "top.bit"))
+            bootloader = artifact_path(args.directory, "software", "bootloader", "bootloader.bin")
 
-        firmwares = []
-        for firmware in "satman", "runtime":
-            filename = artifact_path(args.directory, "software", firmware, firmware + ".fbi")
-            if os.path.exists(filename):
-                firmwares.append(filename)
-        if not firmwares:
-            raise FileNotFoundError("no firmware found")
-        if len(firmwares) > 1:
-            raise ValueError("more than one firmware file, please clean up your build directory. "
-                "Found firmware files: {}".format(" ".join(firmwares)))
-        firmware = firmwares[0]
+            firmwares = []
+            for firmware in "satman", "runtime":
+                filename = artifact_path(args.directory, "software", firmware, firmware + ".fbi")
+                if os.path.exists(filename):
+                    firmwares.append(filename)
+            if not firmwares:
+                raise FileNotFoundError("no firmware found")
+            if len(firmwares) > 1:
+                raise ValueError("more than one firmware file, please clean up your build directory. "
+                    "Found firmware files: {}".format(" ".join(firmwares)))
+            firmware = firmwares[0]
+
+            bins = [ gateware, bootloader, firmware ]
 
-        bins = [ gateware, bootloader, firmware ]
         mgmt.flash(bins)
 
     if args.tool == "reboot":