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 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(&instr)) { + return {alloca, alloca->getAllocatedType()}; + } + if (auto *gep = llvm::dyn_cast(&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()};