From e266fa93ee66e7a4605a82eafdbe7efe2e28e58a Mon Sep 17 00:00:00 2001
From: Alex Cameron <asc@tetsuo.sh>
Date: Tue, 21 Jun 2022 11:32:32 +1000
Subject: [PATCH] Refine APIs

---
 include/remill/Arch/Arch.h |  3 ++-
 lib/Arch/Arch.cpp          | 15 +++++++++------
 2 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/include/remill/Arch/Arch.h b/include/remill/Arch/Arch.h
index 438b9ead..18fb2cc3 100644
--- a/include/remill/Arch/Arch.h
+++ b/include/remill/Arch/Arch.h
@@ -125,7 +125,8 @@ struct Register {
   // The directly enclosed registers.
   std::vector<const Register *> children;
 
-  void ComputeGEPAccessors(const llvm::DataLayout &dl, llvm::Type *state_type);
+  void ComputeGEPAccessors(const llvm::DataLayout &dl,
+                           llvm::StructType *state_type);
 };
 
 class Arch {
diff --git a/lib/Arch/Arch.cpp b/lib/Arch/Arch.cpp
index c75be0e8..052a78d2 100644
--- a/lib/Arch/Arch.cpp
+++ b/lib/Arch/Arch.cpp
@@ -444,7 +444,7 @@ namespace {
 
 // Compute the total offset of a GEP chain.
 static uint64_t TotalOffset(const llvm::DataLayout &dl, llvm::Value *base,
-                            llvm::Type *state_type) {
+                            llvm::StructType *state_type) {
   uint64_t total_offset = 0;
   const auto state_size = dl.getTypeAllocSize(state_type);
   while (base) {
@@ -481,9 +481,8 @@ static uint64_t TotalOffset(const llvm::DataLayout &dl, llvm::Value *base,
 
 static llvm::Value *
 FinishAddressOf(llvm::IRBuilder<> &ir, const llvm::DataLayout &dl,
-                llvm::Type *state_type, size_t state_size,
-                const Register *reg, unsigned addr_space,
-                llvm::Value *gep) {
+                llvm::StructType *state_type, size_t state_size,
+                const Register *reg, unsigned addr_space, llvm::Value *gep) {
 
 
   auto gep_offset = TotalOffset(dl, gep, state_type);
@@ -548,11 +547,15 @@ FinishAddressOf(llvm::IRBuilder<> &ir, const llvm::DataLayout &dl,
 }  // namespace
 
 void Register::ComputeGEPAccessors(const llvm::DataLayout &dl,
-                                   llvm::Type *state_type) {
-  if (gep_type_at_offset || !state_type) {
+                                   llvm::StructType *state_type) {
+  if (gep_type_at_offset) {
     return;
   }
 
+  if (!state_type) {
+    state_type = arch->state_type;
+  }
+
   auto &context = state_type->getContext();
 
   gep_index_list.push_back(