Fix BasicBlockContext lookups

ekilmer · ekilmer · commit 6daea1997ce2 · 2023-09-28T11:52:28.000-04:00
Fixes the missing literal_struct_2 issue
diff --git a/include/anvill/ABI.h b/include/anvill/ABI.h
@@ -86,7 +86,8 @@ extern const std::string kAnvillStackZero;
 // use this to queue off of then just move it after the split
 extern const std::string kStackMetadata;
 
-extern const std::string kBasicBlockMetadata;
+extern const std::string kBasicBlockAddrMetadata;
+extern const std::string kBasicBlockUidMetadata;
 
 
 /// Intrinsic that acts like a return instruction but leaves both the basic block and the parent function.
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>>
-  GetBasicBlockContextForAddr(uint64_t addr) const = 0;
+  GetBasicBlockContextForUid(uint64_t uid) const = 0;
   virtual const FunctionDecl &GetFunctionAtAddress(uint64_t addr) const = 0;
 };
 
@@ -33,9 +33,9 @@ class BasicBlockPass : public llvm::PassInfoMixin<BasicBlockPass<T>> {
   llvm::PreservedAnalyses run(llvm::Function &F,
                               llvm::FunctionAnalysisManager &AM) {
     auto &bb_pass = *static_cast<T *>(this);
-    auto bbaddr = anvill::GetBasicBlockAddr(&F);
-    if (bbaddr.has_value()) {
-      auto maybe_bb_cont = contexts.GetBasicBlockContextForAddr(*bbaddr);
+    auto bbuid = anvill::GetBasicBlockUid(&F);
+    if (bbuid.has_value()) {
+      auto maybe_bb_cont = contexts.GetBasicBlockContextForUid(*bbuid);
       if (maybe_bb_cont) {
         const BasicBlockContext &bb_cont = *maybe_bb_cont;
         auto &parent_func =
diff --git a/include/anvill/Specification.h b/include/anvill/Specification.h
@@ -98,7 +98,7 @@ class SpecBlockContexts : public BasicBlockContexts {
   SpecBlockContexts(const Specification &spec);
 
   virtual std::optional<std::reference_wrapper<const BasicBlockContext>>
-  GetBasicBlockContextForAddr(uint64_t addr) const override;
+  GetBasicBlockContextForUid(uint64_t 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,6 +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);
 
 llvm::Argument *GetBasicBlockStackPtr(llvm::Function *func);
 
diff --git a/lib/ABI.cpp b/lib/ABI.cpp
@@ -82,7 +82,8 @@ const std::string kAnvillDataProvenanceFunc(kAnvillNamePrefix +
 // `alloca`.
 const std::string kAnvillStackZero(kAnvillNamePrefix + "stack_zero");
 
-const std::string kBasicBlockMetadata(kAnvillNamePrefix + "basic_block_md");
+const std::string kBasicBlockAddrMetadata(kAnvillNamePrefix + "basic_block_addr_md");
+const std::string kBasicBlockUidMetadata(kAnvillNamePrefix + "basic_block_uid_md");
 
 const std::string kStackMetadata(kAnvillNamePrefix + "stack_alloc");
 
diff --git a/lib/Declarations.cpp b/lib/Declarations.cpp
@@ -78,8 +78,8 @@ VariableDecl::DeclareInModule(const std::string &name,
 
 void FunctionDecl::AddBBContexts(
     std::unordered_map<uint64_t, SpecBlockContext> &contexts) const {
-  for (const auto &[addr, _] : this->cfg) {
-    contexts.insert({addr, this->GetBlockContext(addr)});
+  for (const auto &[uid, _] : this->cfg) {
+    contexts.insert({uid, this->GetBlockContext(uid)});
   }
 }
 
@@ -475,12 +475,12 @@ void CallableDecl::OverrideFunctionTypeWithABIReturnLayout() {
 
 namespace {
 template <class V>
-V GetWithDef(uint64_t addr, const std::unordered_map<uint64_t, V> &map, V def) {
-  if (map.find(addr) == map.end()) {
+V GetWithDef(uint64_t uid, const std::unordered_map<uint64_t, V> &map, V def) {
+  if (map.find(uid) == map.end()) {
     return def;
   }
 
-  return map.find(addr)->second;
+  return map.find(uid)->second;
 }
 }  // namespace
 
diff --git a/lib/Lifters/BasicBlockLifter.cpp b/lib/Lifters/BasicBlockLifter.cpp
@@ -402,9 +402,12 @@ void BasicBlockLifter::LiftInstructionsIntoLiftedFunction() {
 }
 
 
-llvm::MDNode *BasicBlockLifter::GetBasicBlockAnnotation(uint64_t addr) const {
+llvm::MDNode *BasicBlockLifter::GetBasicBlockAddrAnnotation(uint64_t addr) const {
   return this->GetAddrAnnotation(addr, this->semantics_module->getContext());
 }
+llvm::MDNode *BasicBlockLifter::GetBasicBlockUidAnnotation(uint64_t 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" +
@@ -437,8 +440,10 @@ llvm::Function *BasicBlockLifter::DeclareBasicBlockFunction() {
 
 BasicBlockFunction BasicBlockLifter::CreateBasicBlockFunction() {
   auto func = bb_func;
-  func->setMetadata(anvill::kBasicBlockMetadata,
-                    GetBasicBlockAnnotation(this->block_def.addr));
+  func->setMetadata(anvill::kBasicBlockAddrMetadata,
+                    GetBasicBlockAddrAnnotation(this->block_def.addr));
+  func->setMetadata(anvill::kBasicBlockUidMetadata,
+                    GetBasicBlockUidAnnotation(this->block_def.uid));
 
   auto &context = this->semantics_module->getContext();
   llvm::FunctionType *lifted_func_type =
diff --git a/lib/Lifters/BasicBlockLifter.h b/lib/Lifters/BasicBlockLifter.h
@@ -107,7 +107,8 @@ class BasicBlockLifter : public CodeLifter {
                              remill::DecodingContext context);
 
 
-  llvm::MDNode *GetBasicBlockAnnotation(uint64_t addr) const;
+  llvm::MDNode *GetBasicBlockAddrAnnotation(uint64_t addr) const;
+  llvm::MDNode *GetBasicBlockUidAnnotation(uint64_t uid) const;
 
  public:
   BasicBlockLifter(std::unique_ptr<BasicBlockContext> block_context,
diff --git a/lib/Lifters/CodeLifter.cpp b/lib/Lifters/CodeLifter.cpp
@@ -57,6 +57,8 @@ CodeLifter::CodeLifter(const LifterOptions &options,
       type_specifier(type_specifier),
       address_type(
           llvm::Type::getIntNTy(llvm_context, options.arch->address_size)),
+      uid_type(
+          llvm::Type::getInt64Ty(llvm_context)),
       i8_type(llvm::Type::getInt8Ty(llvm_context)),
       i8_zero(llvm::Constant::getNullValue(i8_type)),
       i32_type(llvm::Type::getInt32Ty(llvm_context)),
@@ -191,6 +193,14 @@ llvm::MDNode *CodeLifter::GetAddrAnnotation(uint64_t addr,
   return llvm::MDNode::get(context, pc_md);
 }
 
+llvm::MDNode *CodeLifter::GetUidAnnotation(uint64_t uid,
+                                            llvm::LLVMContext &context) const {
+  auto uid_val = llvm::ConstantInt::get(
+      remill::RecontextualizeType(uid_type, context), uid);
+  auto uid_md = llvm::ValueAsMetadata::get(uid_val);
+  return llvm::MDNode::get(context, uid_md);
+}
+
 // Allocate and initialize the state structure.
 llvm::Value *
 CodeLifter::AllocateAndInitializeStateStructure(llvm::BasicBlock *block,
diff --git a/lib/Lifters/CodeLifter.h b/lib/Lifters/CodeLifter.h
@@ -45,6 +45,7 @@ class CodeLifter {
   const TypeProvider &type_provider;
   const TypeTranslator &type_specifier;
   llvm::IntegerType *const address_type;
+  llvm::IntegerType *const uid_type;
 
 
   // Convenient to keep around.
@@ -78,6 +79,8 @@ class CodeLifter {
 
   llvm::MDNode *GetAddrAnnotation(uint64_t addr,
                                   llvm::LLVMContext &context) const;
+  llvm::MDNode *GetUidAnnotation(uint64_t uid,
+                                  llvm::LLVMContext &context) const;
 
  public:
   CodeLifter(const LifterOptions &options, llvm::Module *semantics_module,
diff --git a/lib/Lifters/FunctionLifter.cpp b/lib/Lifters/FunctionLifter.cpp
@@ -354,7 +354,6 @@ BasicBlockLifter &FunctionLifter::GetOrCreateBasicBlockLifter(uint64_t uid) {
           this->curr_decl->GetBlockContext(uid));
 
   auto &cfg = this->curr_decl->cfg;
-  CHECK(cfg.contains(uid));
   CodeBlock defblk = cfg.find(uid)->second;
 
   auto inserted = this->bb_lifters.emplace(
@@ -384,8 +383,8 @@ void FunctionLifter::VisitBlocks(llvm::Value *lifted_function_state,
              << ": " << this->curr_decl->cfg.size();
 
 
-  for (const auto &[addr, blk] : this->curr_decl->cfg) {
-    DLOG(INFO) << "Visiting: " << std::hex << addr;
+  for (const auto &[uid, blk] : this->curr_decl->cfg) {
+    DLOG(INFO) << "Visiting: " << std::hex << blk.addr << " " << std::dec << uid;
     this->VisitBlock(blk, lifted_function_state, abstract_stack);
   }
 }
@@ -724,9 +723,13 @@ FunctionLifter::AddFunctionToContext(llvm::Function *func,
       }
       remill::CloneFunctionInto(old_version, new_version);
       new_version->setMetadata(
-          kBasicBlockMetadata,
+          kBasicBlockAddrMetadata,
           this->GetAddrAnnotation(block.addr, module_context));
+      new_version->setMetadata(
+          kBasicBlockUidMetadata,
+          this->GetUidAnnotation(block.uid, module_context));
       CHECK(anvill::GetBasicBlockAddr(new_version).has_value());
+      CHECK(anvill::GetBasicBlockUid(new_version).has_value());
     }
   }
 
diff --git a/lib/Passes/ReplaceStackReferences.cpp b/lib/Passes/ReplaceStackReferences.cpp
@@ -255,7 +255,8 @@ llvm::PreservedAnalyses ReplaceStackReferences::runOnBasicBlockFunction(
       AbstractStack::StackTypeFromSize(F.getContext(), overrunsz));
 
   DLOG(INFO) << "Replacing stack vars in bb: " << std::hex
-             << *anvill::GetBasicBlockAddr(&F);
+             << *anvill::GetBasicBlockAddr(&F) << " " << std::dec
+             << *anvill::GetBasicBlockUid(&F);
   DLOG(INFO) << "Stack size " << cont.GetStackSize();
   DLOG(INFO) << "Max stack size " << cont.GetMaxStackSize();
   AbstractStack stk(
diff --git a/lib/Specification.cpp b/lib/Specification.cpp
@@ -447,8 +447,8 @@ SpecBlockContexts::SpecBlockContexts(const Specification &spec) {
 }
 
 std::optional<std::reference_wrapper<const BasicBlockContext>>
-SpecBlockContexts::GetBasicBlockContextForAddr(uint64_t addr) const {
-  auto cont = this->contexts.find(addr);
+SpecBlockContexts::GetBasicBlockContextForUid(uint64_t 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
@@ -970,7 +970,17 @@ bool CanBeAliased(llvm::Value *val) {
 }
 
 std::optional<uint64_t> GetBasicBlockAddr(llvm::Function *func) {
-  auto meta = func->getMetadata(kBasicBlockMetadata);
+  auto meta = func->getMetadata(kBasicBlockAddrMetadata);
+  if (!meta) {
+    return std::nullopt;
+  }
+
+  auto v = llvm::cast<llvm::ValueAsMetadata>(meta->getOperand(0))->getValue();
+
+  return llvm::cast<llvm::ConstantInt>(v)->getLimitedValue();
+}
+std::optional<uint64_t> GetBasicBlockUid(llvm::Function *func) {
+  auto meta = func->getMetadata(kBasicBlockUidMetadata);
   if (!meta) {
     return std::nullopt;
   }

Original file line number	Diff line number	Diff line change
`@@ -78,8 +78,8 @@ VariableDecl::DeclareInModule(const std::string &name,`
`78`	`78`
`79`	`79`	`void FunctionDecl::AddBBContexts(`
`80`	`80`	`std::unordered_map<uint64_t, SpecBlockContext> &contexts) const {`
`81`		`- for (const auto &[addr, _] : this->cfg) {`
`82`		`- contexts.insert({addr, this->GetBlockContext(addr)});`
	`81`	`+ for (const auto &[uid, _] : this->cfg) {`
	`82`	`+ contexts.insert({uid, this->GetBlockContext(uid)});`
`83`	`83`	`}`
`84`	`84`	`}`
`85`	`85`
`@@ -475,12 +475,12 @@ void CallableDecl::OverrideFunctionTypeWithABIReturnLayout() {`
`475`	`475`
`476`	`476`	`namespace {`
`477`	`477`	`template <class V>`
`478`		`-V GetWithDef(uint64_t addr, const std::unordered_map<uint64_t, V> &map, V def) {`
`479`		`- if (map.find(addr) == map.end()) {`
	`478`	`+V GetWithDef(uint64_t uid, const std::unordered_map<uint64_t, V> &map, V def) {`
	`479`	`+ if (map.find(uid) == map.end()) {`
`480`	`480`	`return def;`
`481`	`481`	`}`
`482`	`482`
`483`		`- return map.find(addr)->second;`
	`483`	`+ return map.find(uid)->second;`
`484`	`484`	`}`
`485`	`485`	`} // namespace`
`486`	`486`