From 3b054855ecb758e6a5697fe4811d38e12cde68f8 Mon Sep 17 00:00:00 2001
From: whitequark <whitequark@whitequark.org>
Date: Sat, 21 Apr 2018 18:27:14 +0000
Subject: [PATCH] firmware: add allocator debug feature, invoked by
 artiq_coredebug frontend.

---
 artiq/coredevice/comm_mgmt.py         |  5 ++++
 artiq/firmware/libproto/mgmt_proto.rs |  3 ++
 artiq/firmware/runtime/mgmt.rs        |  5 +++-
 artiq/frontend/artiq_coredebug.py     | 43 +++++++++++++++++++++++++++
 setup.py                              |  1 +
 5 files changed, 56 insertions(+), 1 deletion(-)
 create mode 100644 artiq/frontend/artiq_coredebug.py

diff --git a/artiq/coredevice/comm_mgmt.py b/artiq/coredevice/comm_mgmt.py
index 9a551a378..f0b7d009a 100644
--- a/artiq/coredevice/comm_mgmt.py
+++ b/artiq/coredevice/comm_mgmt.py
@@ -17,6 +17,8 @@ class Request(Enum):
     Hotswap = 4
     Reboot = 5
 
+    DebugAllocator = 8
+
 
 class Reply(Enum):
     Success = 1
@@ -151,3 +153,6 @@ class CommMgmt:
     def reboot(self):
         self._write_header(Request.Reboot)
         self._read_expect(Reply.RebootImminent)
+
+    def debug_allocator(self):
+        self._write_header(Request.DebugAllocator)
diff --git a/artiq/firmware/libproto/mgmt_proto.rs b/artiq/firmware/libproto/mgmt_proto.rs
index cb5e08856..c96ad27c5 100644
--- a/artiq/firmware/libproto/mgmt_proto.rs
+++ b/artiq/firmware/libproto/mgmt_proto.rs
@@ -16,6 +16,8 @@ pub enum Request {
 
     Hotswap(Vec<u8>),
     Reboot,
+
+    DebugAllocator,
 }
 
 pub enum Reply<'a> {
@@ -52,6 +54,7 @@ impl Request {
             6 => Request::SetUartLogFilter(read_log_level_filter(reader)?),
             4 => Request::Hotswap(reader.read_bytes()?),
             5 => Request::Reboot,
+            8 => Request::DebugAllocator,
             _  => return Err(io::Error::new(io::ErrorKind::InvalidData, "unknown request type"))
         })
     }
diff --git a/artiq/firmware/runtime/mgmt.rs b/artiq/firmware/runtime/mgmt.rs
index 8877f6221..4db85de70 100644
--- a/artiq/firmware/runtime/mgmt.rs
+++ b/artiq/firmware/runtime/mgmt.rs
@@ -97,7 +97,10 @@ fn worker(io: &Io, stream: &mut TcpStream) -> io::Result<()> {
                 stream.close()?;
                 warn!("restarting");
                 unsafe { boot::reset() }
-            }
+            },
+
+            Request::DebugAllocator =>
+                unsafe { println!("{}", ::ALLOC) },
         };
     }
 }
diff --git a/artiq/frontend/artiq_coredebug.py b/artiq/frontend/artiq_coredebug.py
new file mode 100644
index 000000000..55631e037
--- /dev/null
+++ b/artiq/frontend/artiq_coredebug.py
@@ -0,0 +1,43 @@
+#!/usr/bin/env python3
+
+import argparse
+import sys
+import struct
+
+from artiq.tools import verbosity_args, init_logger
+from artiq.master.databases import DeviceDB
+from artiq.coredevice.comm_mgmt import CommMgmt
+
+
+def get_argparser():
+    parser = argparse.ArgumentParser(description="ARTIQ core device debug tool")
+
+    verbosity_args(parser)
+    parser.add_argument("--device-db", default="device_db.py",
+                       help="device database file (default: '%(default)s')")
+
+    subparsers = parser.add_subparsers(dest="action")
+
+    p_allocator = subparsers.add_parser("allocator",
+                                        help="show heap layout")
+
+    return parser
+
+
+def main():
+    args = get_argparser().parse_args()
+    init_logger(args)
+
+    core_addr = DeviceDB(args.device_db).get("core")["arguments"]["host"]
+    mgmt = CommMgmt(core_addr)
+    try:
+        if args.action == "allocator":
+            mgmt.debug_allocator()
+        else:
+            print("An action needs to be specified.", file=sys.stderr)
+            sys.exit(1)
+    finally:
+        mgmt.close()
+
+if __name__ == "__main__":
+    main()
diff --git a/setup.py b/setup.py
index 2cb40bb5c..ac7e9d7e9 100755
--- a/setup.py
+++ b/setup.py
@@ -25,6 +25,7 @@ console_scripts = [
     "artiq_coreconfig = artiq.frontend.artiq_coreconfig:main",
     "artiq_corelog = artiq.frontend.artiq_corelog:main",
     "artiq_coreboot = artiq.frontend.artiq_coreboot:main",
+    "artiq_coredebug = artiq.frontend.artiq_coredebug:main",
     "artiq_ctlmgr = artiq.frontend.artiq_ctlmgr:main",
     "artiq_devtool = artiq.frontend.artiq_devtool:main",
     "artiq_pcap = artiq.frontend.artiq_pcap:main",