From 903ec8243fc8af856ac549fcbfbd78835f2c3556 Mon Sep 17 00:00:00 2001
From: Tim Newsome <tim@sifive.com>
Date: Thu, 28 Sep 2023 11:13:52 -0700
Subject: [PATCH] debug: Add Openocd.set_available()

This helper uses dmi_write commands to mark harts
available/unavailable.
---
 debug/gdbserver.py | 13 ++++++-------
 debug/testlib.py   | 24 ++++++++++++++++++++++++
 2 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/debug/gdbserver.py b/debug/gdbserver.py
index 5a2ba2abe..23515a393 100755
--- a/debug/gdbserver.py
+++ b/debug/gdbserver.py
@@ -1830,8 +1830,7 @@ def test(self):
         # Other hart should have become unavailable.
         if self.target.support_unavailable_control:
             self.server.wait_until_running(self.target.harts)
-            self.server.command(
-                    f"riscv dmi_write 0x1f 0x{(1<<self.hart.id)&0x3:x}")
+            self.server.set_available([self.hart])
         self.gdb.expect(r"\S+ became unavailable.")
         self.gdb.interrupt()
 
@@ -1901,8 +1900,8 @@ def test(self):
         self.gdb.c(wait=False)
         if self.target.support_unavailable_control:
             self.server.wait_until_running([self.hart])
-            self.server.command(
-                    f"riscv dmi_write 0x1f 0x{(~(1<<self.hart.id))&0x3:x}")
+            self.server.set_available(
+                [h for h in self.target.harts if h != self.hart])
         self.gdb.expect(r"\S+ became unavailable.")
         self.gdb.interrupt()
         # gdb might automatically switch to the available hart.
@@ -1933,14 +1932,14 @@ def test(self):
         self.gdb.p("$pc=loop_forever")
         self.gdb.c(wait=False)
         self.server.wait_until_running([self.hart])
-        self.server.command(
-                f"riscv dmi_write 0x1f 0x{(~(1<<self.hart.id))&0x3:x}")
+        self.server.set_available(
+                [h for h in self.target.harts if h != self.hart])
         self.gdb.expect(r"\S+ became unavailable.")
 
         # Now send a DMI command through OpenOCD to make the hart available
         # again.
 
-        self.server.command("riscv dmi_write 0x1f 0x3")
+        self.server.set_available(self.target.harts)
         self.gdb.expect(r"\S+ became available")
         self.gdb.interrupt()
         self.gdb.p("$pc")
diff --git a/debug/testlib.py b/debug/testlib.py
index 435b41b5a..8b8ca58ba 100644
--- a/debug/testlib.py
+++ b/debug/testlib.py
@@ -496,6 +496,30 @@ def wait_until_running(self, harts):
             if time.time() - start > self.timeout:
                 raise TestLibError("Timed out waiting for targets to run.")
 
+    def set_available(self, harts):
+        """Set the given harts to available, and any others to be unavailable.
+        This uses a custom DMI register (0x1f) that is only implemented in
+        spike."""
+        available_mask = 0
+        for hart in harts:
+            available_mask |= 1 << hart.id
+        self.command(f"riscv dmi_write 0x1f 0x{available_mask:x}")
+
+        # Wait until it happened.
+        start = time.time()
+        while True:
+            currently_available = set()
+            currently_unavailable = set()
+            for i, target in enumerate(self.targets()):
+                if target["State"] == "unavailable":
+                    currently_unavailable.add(i)
+                else:
+                    currently_available.add(i)
+            if currently_available == set(hart.id for hart in harts):
+                return
+            if time.time() - start > self.timeout:
+                raise TestLibError("Timed out waiting for hart availability.")
+
 class OpenocdCli:
     def __init__(self, port=4444):
         self.child = pexpect.spawn(