Strong Uid typing

lifting-bits · Sep 29, 2023 · b738508 · b738508
1 parent 4cfe2d9
commit b738508
Show file tree

Hide file tree

Showing 16 changed files with 68 additions and 45 deletions.
diff --git a/include/anvill/Declarations.h b/include/anvill/Declarations.h
@@ -14,6 +14,7 @@
 
 #include <cstddef>
 #include <cstdint>
+#include <functional>
 #include <memory>
 #include <optional>
 #include <string>
@@ -45,16 +46,31 @@ struct Register;
 }  // namespace remill
 namespace anvill {
 
+struct Uid {
+  std::uint64_t value;
+  bool operator==(const Uid &) const = default;
+};
+
+}
+
+template <>
+struct std::hash<anvill::Uid> {
+  size_t operator()(const anvill::Uid &uid) const noexcept {
+    return std::hash<uint64_t>()(uid.value);
+  }
+};
+
+namespace anvill {
 
 struct CodeBlock {
   uint64_t addr;
   uint32_t size;
-  std::unordered_set<uint64_t> outgoing_edges;
+  std::unordered_set<Uid> outgoing_edges;
   // The set of context assignments that occur at the entry point to this block.
   // A block may have specific decoding context properties such as "TM=1" (the thumb bit is set)
   // So we declare the context assignments that occur at the entry point to a block.
   std::unordered_map<std::string, std::uint64_t> context_assignments;
-  uint64_t uid;
+  Uid uid;
 };
 
 
@@ -401,24 +417,24 @@ struct FunctionDecl : public CallableDecl {
   bool is_extern{false};
 
   // These are the blocks contained within the function representing the CFG.
-  std::unordered_map<std::uint64_t, CodeBlock> cfg;
+  std::unordered_map<Uid, CodeBlock> cfg;
 
   std::unordered_map<std::string, ParameterDecl> locals;
 
-  std::unordered_map<std::uint64_t, SpecStackOffsets> stack_offsets_at_entry;
+  std::unordered_map<Uid, SpecStackOffsets> stack_offsets_at_entry;
 
-  std::unordered_map<std::uint64_t, SpecStackOffsets> stack_offsets_at_exit;
+  std::unordered_map<Uid, SpecStackOffsets> stack_offsets_at_exit;
 
-  std::unordered_map<std::uint64_t, std::vector<ParameterDecl>>
+  std::unordered_map<Uid, std::vector<ParameterDecl>>
       live_regs_at_entry;
 
-  std::unordered_map<std::uint64_t, std::vector<ParameterDecl>>
+  std::unordered_map<Uid, std::vector<ParameterDecl>>
       live_regs_at_exit;
 
-  std::unordered_map<std::uint64_t, std::vector<ConstantDomain>>
+  std::unordered_map<Uid, std::vector<ConstantDomain>>
       constant_values_at_entry;
 
-  std::unordered_map<std::uint64_t, std::vector<ConstantDomain>>
+  std::unordered_map<Uid, std::vector<ConstantDomain>>
       constant_values_at_exit;
 
   // sorted vector of hints
@@ -452,10 +468,10 @@ struct FunctionDecl : public CallableDecl {
   static Result<FunctionDecl, std::string> Create(llvm::Function &func,
                                                   const remill::Arch *arch);
 
-  SpecBlockContext GetBlockContext(std::uint64_t uid) const;
+  SpecBlockContext GetBlockContext(Uid uid) const;
 
   void
-  AddBBContexts(std::unordered_map<uint64_t, SpecBlockContext> &contexts) const;
+  AddBBContexts(std::unordered_map<Uid, SpecBlockContext> &contexts) const;
 };
 
 // A call site decl, as represented at a "near ABI" level. This is like a

diff --git a/include/anvill/Passes/BasicBlockPass.h b/include/anvill/Passes/BasicBlockPass.h
@@ -15,7 +15,7 @@ namespace anvill {
 class BasicBlockContexts {
  public:
   virtual std::optional<std::reference_wrapper<const BasicBlockContext>>
-  GetBasicBlockContextForUid(uint64_t uid) const = 0;
+  GetBasicBlockContextForUid(Uid uid) const = 0;
   virtual const FunctionDecl &GetFunctionAtAddress(uint64_t addr) const = 0;
 };
 

diff --git a/include/anvill/Specification.h b/include/anvill/Specification.h
@@ -91,14 +91,14 @@ struct ValueDecl;
 
 class Specification;
 class SpecBlockContexts : public BasicBlockContexts {
-  std::unordered_map<uint64_t, SpecBlockContext> contexts;
+  std::unordered_map<Uid, SpecBlockContext> contexts;
   std::unordered_map<uint64_t, std::shared_ptr<const FunctionDecl>> funcs;
 
  public:
   SpecBlockContexts(const Specification &spec);
 
   virtual std::optional<std::reference_wrapper<const BasicBlockContext>>
-  GetBasicBlockContextForUid(uint64_t uid) const override;
+  GetBasicBlockContextForUid(Uid uid) const override;
 
   virtual const FunctionDecl &
   GetFunctionAtAddress(uint64_t addr) const override;

diff --git a/include/anvill/Utils.h b/include/anvill/Utils.h
@@ -135,7 +135,7 @@ llvm::Value *StoreNativeValue(llvm::Value *native_val, const ValueDecl &decl,
                               llvm::Value *state_ptr, llvm::Value *mem_ptr);
 
 std::optional<uint64_t> GetBasicBlockAddr(llvm::Function *func);
-std::optional<uint64_t> GetBasicBlockUid(llvm::Function *func);
+std::optional<Uid> GetBasicBlockUid(llvm::Function *func);
 
 llvm::Argument *GetBasicBlockStackPtr(llvm::Function *func);
 

diff --git a/lib/Declarations.cpp b/lib/Declarations.cpp
@@ -77,7 +77,7 @@ VariableDecl::DeclareInModule(const std::string &name,
 }
 
 void FunctionDecl::AddBBContexts(
-    std::unordered_map<uint64_t, SpecBlockContext> &contexts) const {
+    std::unordered_map<Uid, SpecBlockContext> &contexts) const {
   for (const auto &[uid, _] : this->cfg) {
     contexts.insert({uid, this->GetBlockContext(uid)});
   }
@@ -475,7 +475,7 @@ void CallableDecl::OverrideFunctionTypeWithABIReturnLayout() {
 
 namespace {
 template <class V>
-V GetWithDef(uint64_t uid, const std::unordered_map<uint64_t, V> &map, V def) {
+V GetWithDef(Uid uid, const std::unordered_map<Uid, V> &map, V def) {
   if (map.find(uid) == map.end()) {
     return def;
   }
@@ -488,7 +488,7 @@ size_t FunctionDecl::GetPointerDisplacement() const {
   return this->parameter_size + this->parameter_offset;
 }
 
-SpecBlockContext FunctionDecl::GetBlockContext(std::uint64_t uid) const {
+SpecBlockContext FunctionDecl::GetBlockContext(Uid uid) const {
   return SpecBlockContext(
       *this, GetWithDef(uid, this->stack_offsets_at_entry, SpecStackOffsets()),
       GetWithDef(uid, this->stack_offsets_at_exit, SpecStackOffsets()),

diff --git a/lib/Lifters/BasicBlockLifter.cpp b/lib/Lifters/BasicBlockLifter.cpp
@@ -405,13 +405,13 @@ void BasicBlockLifter::LiftInstructionsIntoLiftedFunction() {
 llvm::MDNode *BasicBlockLifter::GetBasicBlockAddrAnnotation(uint64_t addr) const {
   return this->GetAddrAnnotation(addr, this->semantics_module->getContext());
 }
-llvm::MDNode *BasicBlockLifter::GetBasicBlockUidAnnotation(uint64_t uid) const {
+llvm::MDNode *BasicBlockLifter::GetBasicBlockUidAnnotation(Uid uid) const {
   return this->GetUidAnnotation(uid, this->semantics_module->getContext());
 }
 
 llvm::Function *BasicBlockLifter::DeclareBasicBlockFunction() {
   std::string name_ = "func" + std::to_string(decl.address) + "basic_block" +
-                      std::to_string(this->block_def.addr) + "_" + std::to_string(this->block_def.uid);
+                      std::to_string(this->block_def.addr) + "_" + std::to_string(this->block_def.uid.value);
   auto &context = this->semantics_module->getContext();
   llvm::FunctionType *lifted_func_type =
       llvm::dyn_cast<llvm::FunctionType>(remill::RecontextualizeType(

diff --git a/lib/Lifters/BasicBlockLifter.h b/lib/Lifters/BasicBlockLifter.h
@@ -108,7 +108,7 @@ class BasicBlockLifter : public CodeLifter {
 
 
   llvm::MDNode *GetBasicBlockAddrAnnotation(uint64_t addr) const;
-  llvm::MDNode *GetBasicBlockUidAnnotation(uint64_t uid) const;
+  llvm::MDNode *GetBasicBlockUidAnnotation(Uid uid) const;
 
  public:
   BasicBlockLifter(std::unique_ptr<BasicBlockContext> block_context,

diff --git a/lib/Lifters/CodeLifter.cpp b/lib/Lifters/CodeLifter.cpp
@@ -20,6 +20,8 @@
 
 #include <unordered_set>
 
+#include "anvill/Declarations.h"
+
 namespace anvill {
 namespace {
 // Clear out LLVM variable names. They're usually not helpful.
@@ -193,10 +195,10 @@ llvm::MDNode *CodeLifter::GetAddrAnnotation(uint64_t addr,
   return llvm::MDNode::get(context, pc_md);
 }
 
-llvm::MDNode *CodeLifter::GetUidAnnotation(uint64_t uid,
+llvm::MDNode *CodeLifter::GetUidAnnotation(Uid uid,
                                             llvm::LLVMContext &context) const {
   auto uid_val = llvm::ConstantInt::get(
-      remill::RecontextualizeType(uid_type, context), uid);
+      remill::RecontextualizeType(uid_type, context), uid.value);
   auto uid_md = llvm::ValueAsMetadata::get(uid_val);
   return llvm::MDNode::get(context, uid_md);
 }

diff --git a/lib/Lifters/CodeLifter.h b/lib/Lifters/CodeLifter.h
@@ -7,6 +7,7 @@
 #include <remill/BC/InstructionLifter.h>
 #include <remill/BC/IntrinsicTable.h>
 
+#include "anvill/Declarations.h"
 #include "anvill/Lifters.h"
 
 namespace anvill {
@@ -79,7 +80,7 @@ class CodeLifter {
 
   llvm::MDNode *GetAddrAnnotation(uint64_t addr,
                                   llvm::LLVMContext &context) const;
-  llvm::MDNode *GetUidAnnotation(uint64_t uid,
+  llvm::MDNode *GetUidAnnotation(Uid uid,
                                   llvm::LLVMContext &context) const;
 
  public:

diff --git a/lib/Lifters/FunctionLifter.cpp b/lib/Lifters/FunctionLifter.cpp
@@ -336,15 +336,15 @@ llvm::Function *FunctionLifter::DeclareFunction(const FunctionDecl &decl) {
   return GetOrDeclareFunction(decl);
 }
 
-static uint64_t GetRandUid() {
+static Uid GetRandUid() {
   static std::random_device rd;
   static std::mt19937_64 engine(rd());
   static std::uniform_int_distribution<uint64_t> dist(0, UINT64_MAX);
-  return dist(engine);
+  return {dist(engine)};
 }
 
-BasicBlockLifter &FunctionLifter::GetOrCreateBasicBlockLifter(uint64_t uid) {
-  std::pair<uint64_t, uint64_t> key{curr_decl->address, uid};
+BasicBlockLifter &FunctionLifter::GetOrCreateBasicBlockLifter(Uid uid) {
+  std::pair<uint64_t, uint64_t> key{curr_decl->address, uid.value};
   auto lifter = this->bb_lifters.find(key);
   if (lifter != this->bb_lifters.end()) {
     return lifter->second;
@@ -384,7 +384,7 @@ void FunctionLifter::VisitBlocks(llvm::Value *lifted_function_state,
 
 
   for (const auto &[uid, blk] : this->curr_decl->cfg) {
-    DLOG(INFO) << "Visiting: " << std::hex << blk.addr << " " << std::dec << uid;
+    DLOG(INFO) << "Visiting: " << std::hex << blk.addr << " " << std::dec << uid.value;
     this->VisitBlock(blk, lifted_function_state, abstract_stack);
   }
 }
@@ -710,7 +710,7 @@ FunctionLifter::AddFunctionToContext(llvm::Function *func,
   if (!func->isDeclaration()) {
     for (auto &[block_uid, block] : decl.cfg) {
       CHECK(block_uid == block.uid);
-      std::string name = prefix + "basic_block" + std::to_string(block.addr) + "_" + std::to_string(block.uid);
+      std::string name = prefix + "basic_block" + std::to_string(block.addr) + "_" + std::to_string(block.uid.value);
 
       auto new_version = target_module->getFunction(name);
       auto old_version = semantics_module->getFunction(name);

diff --git a/lib/Lifters/FunctionLifter.h b/lib/Lifters/FunctionLifter.h
@@ -106,7 +106,7 @@ class FunctionLifter : public CodeLifter {
 
   // Get or create a basic block lifter for the basic block with specified
   // uid. If a lifter for the uid does not exist, this function will create it
-  BasicBlockLifter &GetOrCreateBasicBlockLifter(uint64_t uid);
+  BasicBlockLifter &GetOrCreateBasicBlockLifter(Uid uid);
 
   const BasicBlockLifter &LiftBasicBlockFunction(const CodeBlock &);
 

diff --git a/lib/Passes/ReplaceStackReferences.cpp b/lib/Passes/ReplaceStackReferences.cpp
@@ -256,7 +256,7 @@ llvm::PreservedAnalyses ReplaceStackReferences::runOnBasicBlockFunction(
 
   DLOG(INFO) << "Replacing stack vars in bb: " << std::hex
              << *anvill::GetBasicBlockAddr(&F) << " " << std::dec
-             << *anvill::GetBasicBlockUid(&F);
+             << (*anvill::GetBasicBlockUid(&F)).value;
   DLOG(INFO) << "Stack size " << cont.GetStackSize();
   DLOG(INFO) << "Max stack size " << cont.GetMaxStackSize();
   AbstractStack stk(

diff --git a/lib/Protobuf.cpp b/lib/Protobuf.cpp
@@ -611,11 +611,11 @@ Result<FunctionDecl, std::string> ProtobufTranslator::DecodeFunction(
 }
 
 void ProtobufTranslator::AddLiveValuesToBB(
-    std::unordered_map<uint64_t, std::vector<ParameterDecl>> &map,
-    uint64_t bb_addr,
+    std::unordered_map<Uid, std::vector<ParameterDecl>> &map,
+    Uid bb_uid,
     const ::google::protobuf::RepeatedPtrField<::specification::Parameter>
         &values) const {
-  auto &v = map.insert({bb_addr, std::vector<ParameterDecl>()}).first->second;
+  auto &v = map.insert({bb_uid, std::vector<ParameterDecl>()}).first->second;
 
   for (auto var : values) {
     auto param = DecodeParameter(var);
@@ -630,23 +630,27 @@ void ProtobufTranslator::AddLiveValuesToBB(
 void ProtobufTranslator::ParseCFGIntoFunction(
     const ::specification::Function &obj, FunctionDecl &decl) const {
   for (const auto& blk : obj.blocks()) {
+    std::unordered_set<Uid> tmp;
+    for (auto o : blk.second.outgoing_blocks()) {
+      tmp.insert({o});
+    }
     CodeBlock nblk = {
         blk.second.address(),
         blk.second.size(),
-        {blk.second.outgoing_blocks().begin(),
-         blk.second.outgoing_blocks().end()},
+        tmp,
         {blk.second.context_assignments().begin(),
          blk.second.context_assignments().end()},
-         blk.first,
+      {blk.first},
     };
-    decl.cfg.emplace(blk.first, std::move(nblk));
+    decl.cfg.emplace(Uid{blk.first}, std::move(nblk));
   }
 
 
-  for (auto &[blk_uid, ctx] : obj.block_context()) {
+  for (auto &[blk_uid_, ctx] : obj.block_context()) {
     std::vector<OffsetDomain> stack_offsets_at_entry, stack_offsets_at_exit;
     std::vector<ConstantDomain> constant_values_at_entry,
         constant_values_at_exit;
+    Uid blk_uid = {blk_uid_};
     auto blk = decl.cfg[blk_uid];
     auto symval_to_domains = [&](const specification::ValueMapping &symval,
                                  std::vector<OffsetDomain> &stack_offsets,

diff --git a/lib/Protobuf.h b/lib/Protobuf.h
@@ -90,8 +90,8 @@ class ProtobufTranslator {
                             FunctionDecl &decl) const;
 
   void AddLiveValuesToBB(
-      std::unordered_map<uint64_t, std::vector<ParameterDecl>> &map,
-      uint64_t bb_addr,
+      std::unordered_map<Uid, std::vector<ParameterDecl>> &map,
+      Uid bb_addr,
       const ::google::protobuf::RepeatedPtrField<::specification::Parameter>
           &values) const;
 

diff --git a/lib/Specification.cpp b/lib/Specification.cpp
@@ -447,7 +447,7 @@ SpecBlockContexts::SpecBlockContexts(const Specification &spec) {
 }
 
 std::optional<std::reference_wrapper<const BasicBlockContext>>
-SpecBlockContexts::GetBasicBlockContextForUid(uint64_t uid) const {
+SpecBlockContexts::GetBasicBlockContextForUid(Uid uid) const {
   auto cont = this->contexts.find(uid);
   if (cont == this->contexts.end()) {
     return std::nullopt;

diff --git a/lib/Utils.cpp b/lib/Utils.cpp
@@ -979,15 +979,15 @@ std::optional<uint64_t> GetBasicBlockAddr(llvm::Function *func) {
 
   return llvm::cast<llvm::ConstantInt>(v)->getLimitedValue();
 }
-std::optional<uint64_t> GetBasicBlockUid(llvm::Function *func) {
+std::optional<Uid> GetBasicBlockUid(llvm::Function *func) {
   auto meta = func->getMetadata(kBasicBlockUidMetadata);
   if (!meta) {
     return std::nullopt;
   }
 
   auto v = llvm::cast<llvm::ValueAsMetadata>(meta->getOperand(0))->getValue();
 
-  return llvm::cast<llvm::ConstantInt>(v)->getLimitedValue();
+  return Uid{llvm::cast<llvm::ConstantInt>(v)->getLimitedValue()};
 }
 
 llvm::Argument *GetBasicBlockStackPtr(llvm::Function *func) {