From 669166d5725e51dc045efbe6ae05fd6d61262175 Mon Sep 17 00:00:00 2001
From: Alex Cameron <asc@tetsuo.sh>
Date: Thu, 23 Jun 2022 22:11:32 +1000
Subject: [PATCH] Retrieve types properly from missing cases in
 `FindVarInFunction`

---
 lib/BC/Util.cpp | 14 ++++++--------
 1 file changed, 6 insertions(+), 8 deletions(-)

diff --git a/lib/BC/Util.cpp b/lib/BC/Util.cpp
index 47def6b5d..c4245553d 100644
--- a/lib/BC/Util.cpp
+++ b/lib/BC/Util.cpp
@@ -218,22 +218,20 @@ FindVarInFunction(llvm::BasicBlock *block, std::string_view name,
 std::pair<llvm::Value *, llvm::Type *>
 FindVarInFunction(llvm::Function *function, std::string_view name_,
                   bool allow_failure) {
-  // TODO(alex): Figure out how to to get types for the two cases below.
   llvm::StringRef name(name_.data(), name_.size());
   if (!function->empty()) {
     for (auto &instr : function->getEntryBlock()) {
       if (instr.getName() == name) {
-        return {&instr, nullptr};
+        if (auto *alloca = llvm::dyn_cast<llvm::AllocaInst>(&instr)) {
+          return {alloca, alloca->getAllocatedType()};
+        }
+        if (auto *gep = llvm::dyn_cast<llvm::GetElementPtrInst>(&instr)) {
+          return {gep, gep->getResultElementType()};
+        }
       }
     }
   }
 
-  for (auto &arg : function->args()) {
-    if (arg.getName() == name) {
-      return {&arg, nullptr};
-    }
-  }
-
   auto module = function->getParent();
   if (auto var = module->getGlobalVariable(name)) {
     return {var, var->getValueType()};