More aarch64 fixes

lifting-bits · Sep 23, 2020 · 3b2f44f · 3b2f44f
1 parent 590c9b8
commit 3b2f44f
Show file tree

Hide file tree

Showing 3 changed files with 56 additions and 22 deletions.
diff --git a/remill/Arch/AArch64/Arch.cpp b/remill/Arch/AArch64/Arch.cpp
@@ -1220,6 +1220,17 @@ bool AArch64Arch::DecodeInstruction(uint64_t address,
     return false;
   }
 
+  // Control flow operands update the next program counter.
+  if (inst.IsControlFlow()) {
+    inst.operands.emplace_back();
+    auto &dst_ret_pc = inst.operands.back();
+    dst_ret_pc.type = Operand::kTypeRegister;
+    dst_ret_pc.action = Operand::kActionWrite;
+    dst_ret_pc.size = address_size;
+    dst_ret_pc.reg.name = "NEXT_PC";
+    dst_ret_pc.reg.size = address_size;
+  }
+
   // The semantics will store the return address in `RETURN_PC`. This is to
   // help synchronize program counters when lifting instructions on an ISA
   // with delay slots.
@@ -1228,9 +1239,9 @@ bool AArch64Arch::DecodeInstruction(uint64_t address,
     auto &dst_ret_pc = inst.operands.back();
     dst_ret_pc.type = Operand::kTypeRegister;
     dst_ret_pc.action = Operand::kActionWrite;
-    dst_ret_pc.size = 64;
+    dst_ret_pc.size = address_size;
     dst_ret_pc.reg.name = "RETURN_PC";
-    dst_ret_pc.reg.size = 64;
+    dst_ret_pc.reg.size = address_size;
   }
 
   return true;

diff --git a/remill/Arch/AArch64/Semantics/BRANCH.cpp b/remill/Arch/AArch64/Semantics/BRANCH.cpp
@@ -117,71 +117,90 @@ DEF_COND(AL) = CondAL;
 
 namespace {
 
-DEF_SEM(DoDirectBranch, PC target_pc) {
-  Write(REG_PC, Read(target_pc));
+DEF_SEM(DoDirectBranch, PC target_pc, R64W pc_dst) {
+  const auto new_pc = Read(target_pc);
+  Write(REG_PC, new_pc);
+  Write(pc_dst, new_pc);
   return memory;
 }
 
 template <typename S>
-DEF_SEM(DoIndirectBranch, S dst) {
-  Write(REG_PC, Read(dst));
+DEF_SEM(DoIndirectBranch, S dst, R64W pc_dst) {
+  const auto new_pc = Read(dst);
+  Write(REG_PC, new_pc);
+  Write(pc_dst, new_pc);
   return memory;
 }
 
 template <bool (*check_cond)(const State &)>
-DEF_SEM(DirectCondBranch, R8W cond, PC taken, PC not_taken) {
+DEF_SEM(DirectCondBranch, R8W cond, PC taken, PC not_taken, R64W pc_dst) {
   addr_t taken_pc = Read(taken);
   addr_t not_taken_pc = Read(not_taken);
   uint8_t take_branch = check_cond(state);
   Write(cond, take_branch);
-  Write(REG_PC, Select<addr_t>(take_branch, taken_pc, not_taken_pc));
+
+  const auto new_pc = Select<addr_t>(take_branch, taken_pc, not_taken_pc);
+  Write(REG_PC, new_pc);
+  Write(pc_dst, new_pc);
   return memory;
 }
 
 template <typename S>
-DEF_SEM(CBZ, R8W cond, PC taken, PC not_taken, S src) {
+DEF_SEM(CBZ, R8W cond, PC taken, PC not_taken, S src, R64W pc_dst) {
   addr_t taken_pc = Read(taken);
   addr_t not_taken_pc = Read(not_taken);
   uint8_t take_branch = UCmpEq(Read(src), 0);
   Write(cond, take_branch);
-  Write(REG_PC, Select<addr_t>(take_branch, taken_pc, not_taken_pc));
+
+  const auto new_pc = Select<addr_t>(take_branch, taken_pc, not_taken_pc);
+  Write(REG_PC, new_pc);
+  Write(pc_dst, new_pc);
   return memory;
 }
 
 template <typename S>
-DEF_SEM(CBNZ, R8W cond, PC taken, PC not_taken, S src) {
+DEF_SEM(CBNZ, R8W cond, PC taken, PC not_taken, S src, R64W pc_dst) {
   addr_t taken_pc = Read(taken);
   addr_t not_taken_pc = Read(not_taken);
   uint8_t take_branch = UCmpNeq(Read(src), 0);
   Write(cond, take_branch);
-  Write(REG_PC, Select<addr_t>(take_branch, taken_pc, not_taken_pc));
+
+  const auto new_pc = Select<addr_t>(take_branch, taken_pc, not_taken_pc);
+  Write(REG_PC, new_pc);
+  Write(pc_dst, new_pc);
   return memory;
 }
 
 
 template <typename S>
-DEF_SEM(TBZ, I8 bit_pos, R8W cond, PC taken, PC not_taken, S src) {
+DEF_SEM(TBZ, I8 bit_pos, R8W cond, PC taken, PC not_taken, S src, R64W pc_dst) {
   addr_t taken_pc = Read(taken);
   addr_t not_taken_pc = Read(not_taken);
   auto bit_n = ZExtTo<S>(Read(bit_pos));
   auto reg_val = ZExtTo<S>(Read(src));
   auto bit_set = UAnd(reg_val, UShl(ZExtTo<S>(1), bit_n));
   auto take_branch = UCmpEq(bit_set, 0);
   Write(cond, take_branch);
-  Write(REG_PC, Select<addr_t>(take_branch, taken_pc, not_taken_pc));
+
+  const auto new_pc = Select<addr_t>(take_branch, taken_pc, not_taken_pc);
+  Write(REG_PC, new_pc);
+  Write(pc_dst, new_pc);
   return memory;
 }
 
 template <typename S>
-DEF_SEM(TBNZ, I8 bit_pos, R8W cond, PC taken, PC not_taken, S src) {
+DEF_SEM(TBNZ, I8 bit_pos, R8W cond, PC taken, PC not_taken, S src, R64W pc_dst) {
   addr_t taken_pc = Read(taken);
   addr_t not_taken_pc = Read(not_taken);
   auto bit_n = ZExtTo<S>(Read(bit_pos));
   auto reg_val = ZExtTo<S>(Read(src));
   auto bit_set = UAnd(reg_val, UShl(ZExtTo<S>(1), bit_n));
   auto take_branch = UCmpNeq(bit_set, 0);
   Write(cond, take_branch);
-  Write(REG_PC, Select<addr_t>(take_branch, taken_pc, not_taken_pc));
+
+  const auto new_pc = Select<addr_t>(take_branch, taken_pc, not_taken_pc);
+  Write(REG_PC, new_pc);
+  Write(pc_dst, new_pc);
   return memory;
 }
 

diff --git a/remill/Arch/AArch64/Semantics/CALL_RET.cpp b/remill/Arch/AArch64/Semantics/CALL_RET.cpp
@@ -17,16 +17,20 @@
 namespace {
 
 template <typename S>
-DEF_SEM(CALL, S target_addr, PC ret_addr, R64W return_pc_dst) {
-  Write(REG_LP, Read(ret_addr));
-  const auto return_pc = Read(target_addr);
-  Write(REG_PC, return_pc);
+DEF_SEM(CALL, S target_addr, PC ret_addr, R64W dst_pc, R64W return_pc_dst) {
+  const auto return_pc = Read(ret_addr);
+  const auto new_pc = Read(target_addr);
+  Write(REG_LP, return_pc);
+  Write(REG_PC, new_pc);
+  Write(dst_pc, new_pc);
   Write(return_pc_dst, return_pc);
   return memory;
 }
 
-DEF_SEM(RET, R64 target_pc) {
-  Write(REG_PC, Read(target_pc));
+DEF_SEM(RET, R64 target_pc, R64W dst_pc) {
+  const auto new_pc = Read(target_pc);
+  Write(REG_PC, new_pc);
+  Write(dst_pc, new_pc);
   return memory;
 }