[lldb][AArch64] Save/restore TLS registers around expressions

[DavidSpickett]

Summary

Previously lldb was storing them but not restoring them. Meaning that this function:

void expr(uint64_t value) {
  __asm__ volatile("msr tpidr_el0, %0" ::"r"(value));
}

When run from lldb:

(lldb) expression expr()

Would leave tpidr as value instead of the original value of the register.

A check for this scenario has been added to TestAArch64LinuxTLSRegisters.py,
which covers tpidr and the SME excluisve tpidr2 register when it's present.

[DavidSpickett]

rebase

[DavidSpickett]

ping!

[omjavaid]

This looks good. Just a nit the summary seems to be using a lot pronouns. Could you kindly improve summary a little and add a one liner about expr_func and how it tests save/restore for future readers of this log.

lldb/source/Plugins/Process/Linux/NativeRegisterContextLinux_arm64.cpp

 //===-- NativeRegisterContextLinux_arm64.cpp ------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
if (GetRegisterInfo().IsSVEEnabled() || GetRegisterInfo().IsSSVEEnabled()) {
     dst += sizeof(m_sve_state);
     ::memcpy(dst, GetSVEBuffer(), GetSVEBufferSize());
     dst += GetSVEBufferSize();
   } else {
     ::memcpy(dst, GetFPRBuffer(), GetFPRSize());
     dst += GetFPRSize();
   }
 
-  if (GetRegisterInfo().IsMTEEnabled())
+  if (GetRegisterInfo().IsMTEEnabled()) {
     ::memcpy(dst, GetMTEControl(), GetMTEControlSize());
+    dst += GetMTEControlSize();
+  }
 
   ::memcpy(dst, GetTLSBuffer(), GetTLSBufferSize());
 
   return error;
 }
 
 Status NativeRegisterContextLinux_arm64::WriteAllRegisterValues(
     const lldb::DataBufferSP &data_sp) {
Status NativeRegisterContextLinux_arm64::WriteAllRegisterValues(
   if (error.Fail())
     return error;
 
   if (GetRegisterInfo().IsMTEEnabled() &&
       data_sp->GetByteSize() > reg_data_min_size) {
     ::memcpy(GetMTEControl(), src, GetMTEControlSize());
     m_mte_ctrl_is_valid = true;
     error = WriteMTEControl();
+    if (error.Fail())
+      return error;
+    src += GetMTEControlSize();
   }
 
+  // There is always a TLS set. It changes size based on system properties, it's
+  // not something an expression can change.
+  ::memcpy(GetTLSBuffer(), src, GetTLSBufferSize());
+  m_tls_is_valid = true;
+  error = WriteTLS();
+
   return error;
 }
 
 bool NativeRegisterContextLinux_arm64::IsGPR(unsigned reg) const {
   if (GetRegisterInfo().GetRegisterSetFromRegisterIndex(reg) ==
       RegisterInfoPOSIX_arm64::GPRegSet)
     return true;
   return false;

lldb/test/API/linux/aarch64/tls_registers/TestAArch64LinuxTLSRegisters.py

def setup(self, registers):
             self.fail("Test program failed to run.")
 
         self.expect(
             "thread list",
             STOPPED_DUE_TO_BREAKPOINT,
             substrs=["stopped", "stop reason = breakpoint"],
         )
 
-    def check_tls_reg(self, registers):
-        self.setup(registers)
-
+    def check_registers(self, registers, values):
         regs = self.thread().GetSelectedFrame().GetRegisters()
         tls_regs = regs.GetFirstValueByName("Thread Local Storage Registers")
         self.assertTrue(tls_regs.IsValid(), "No TLS registers found.")
 
+        for register in registers:
+            tls_reg = tls_regs.GetChildMemberWithName(register)
+            self.assertTrue(tls_reg.IsValid(), "{} register not found.".format(
+                register))
+            self.assertEqual(tls_reg.GetValueAsUnsigned(), values[register])
+
+    def check_tls_reg(self, registers):
+        self.setup(registers)
+
         # Since we can't predict what the value will be, the program has set
         # a target value for us to find.
         initial_values = {
             "tpidr": 0x1122334455667788,
             "tpidr2": 0x8877665544332211,
         }
 
-        for register in registers:
-            tls_reg = tls_regs.GetChildMemberWithName(register)
-            self.assertTrue(tls_reg.IsValid(), "{} register not found.".format(
-                register))
-            self.assertEqual(tls_reg.GetValueAsUnsigned(), initial_values[register])
+        self.check_registers(registers, initial_values)
+
+        # Their values should be restored if an expression modifies them.
+        self.runCmd("expression expr_func()")
+
+        self.check_registers(registers, initial_values)
 
         set_values = {
             "tpidr": 0x1111222233334444,
             "tpidr2": 0x4444333322221111,
         }
 
         # Set our own value(s) for the program to find.
         for register in registers:
def test_tls_no_sme(self):
             self.skipTest("SME must not be present.")
 
         self.check_tls_reg(["tpidr"])
 
     @skipUnlessArch("aarch64")
     @skipUnlessPlatform(["linux"])
     def test_tls_sme(self):
         if not self.isAArch64SME():
-            self.skipTest("SME must present.")
+            self.skipTest("SME must be present.")
 
         self.check_tls_reg(["tpidr", "tpidr2"])
 
     @skipUnlessArch("aarch64")
     @skipUnlessPlatform(["linux"])
     def test_tpidr2_no_sme(self):
         if self.isAArch64SME():
             self.skipTest("SME must not be present.")

lldb/test/API/linux/aarch64/tls_registers/main.c

 void set_tpidr(uint64_t value) {
   __asm__ volatile("msr tpidr_el0, %0" ::"r"(value));
 }
 
 void set_tpidr2(uint64_t value) {
   __asm__ volatile("msr S3_3_C13_C0_5, %0" ::"r"(value));
 }
 
+bool use_tpidr2 = false;
+const uint64_t tpidr_pattern = 0x1122334455667788;
+const uint64_t tpidr2_pattern = 0x8877665544332211;
+
+void expr_func() {
+  set_tpidr(~tpidr_pattern);
+  if (use_tpidr2)
+    set_tpidr2(~tpidr2_pattern);
+}
+
 int main(int argc, char *argv[]) {
-  bool use_tpidr2 = argc > 1;
+  use_tpidr2 = argc > 1;
 
   uint64_t original_tpidr = get_tpidr();
   // Accessing this on a core without it produces SIGILL. Only do this if
   // requested.
   uint64_t original_tpidr2 = 0;
   if (use_tpidr2)
     original_tpidr2 = get_tpidr2();
 
-  uint64_t tpidr_pattern = 0x1122334455667788;
   set_tpidr(tpidr_pattern);
 
-  uint64_t tpidr2_pattern = 0x8877665544332211;
   if (use_tpidr2)
     set_tpidr2(tpidr2_pattern);
 
   // Set break point at this line.
   // lldb will now set its own pattern(s) for us to find.
 
   uint64_t new_tpidr = get_tpidr();
   volatile bool tpidr_was_set = new_tpidr == 0x1111222233334444;