From 8ada962a34cc40241670cb490cb7ed8ce3b98b30 Mon Sep 17 00:00:00 2001
From: Florian Mayer <fmayer@google.com>
Date: Wed, 2 Feb 2022 17:07:19 -0800
Subject: [PATCH] [NFC] [hwasan] use InstIterator

Differential Revision: https://reviews.llvm.org/D118865
---
 .../Instrumentation/HWAddressSanitizer.cpp         | 100 ++++++++++-----------
 1 file changed, 48 insertions(+), 52 deletions(-)

diff --git a/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp b/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp
index 357225d..b7f6487 100644
--- a/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp
@@ -33,6 +33,7 @@
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InlineAsm.h"
+#include "llvm/IR/InstIterator.h"
 #include "llvm/IR/InstVisitor.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
@@ -1480,59 +1481,56 @@ bool HWAddressSanitizer::sanitizeFunction(
   SmallVector<Instruction *, 4> UnrecognizedLifetimes;
   DenseMap<AllocaInst *, std::vector<DbgVariableIntrinsic *>> AllocaDbgMap;
   bool CallsReturnTwice = false;
-  for (auto &BB : F) {
-    for (auto &Inst : BB) {
-      if (CallInst *CI = dyn_cast<CallInst>(&Inst)) {
-        if (CI->canReturnTwice()) {
-          CallsReturnTwice = true;
-        }
+  for (auto &Inst : instructions(F)) {
+    if (CallInst *CI = dyn_cast<CallInst>(&Inst)) {
+      if (CI->canReturnTwice()) {
+        CallsReturnTwice = true;
+      }
+    }
+    if (InstrumentStack) {
+      if (AllocaInst *AI = dyn_cast<AllocaInst>(&Inst)) {
+        if (isInterestingAlloca(*AI))
+          AllocasToInstrument.insert({AI, {}});
+        continue;
       }
-      if (InstrumentStack) {
-        if (AllocaInst *AI = dyn_cast<AllocaInst>(&Inst)) {
-          if (isInterestingAlloca(*AI))
-            AllocasToInstrument.insert({AI, {}});
+      auto *II = dyn_cast<IntrinsicInst>(&Inst);
+      if (II && (II->getIntrinsicID() == Intrinsic::lifetime_start ||
+                 II->getIntrinsicID() == Intrinsic::lifetime_end)) {
+        AllocaInst *AI = findAllocaForValue(II->getArgOperand(1));
+        if (!AI) {
+          UnrecognizedLifetimes.push_back(&Inst);
           continue;
         }
-        auto *II = dyn_cast<IntrinsicInst>(&Inst);
-        if (II && (II->getIntrinsicID() == Intrinsic::lifetime_start ||
-                   II->getIntrinsicID() == Intrinsic::lifetime_end)) {
-          AllocaInst *AI = findAllocaForValue(II->getArgOperand(1));
-          if (!AI) {
-            UnrecognizedLifetimes.push_back(&Inst);
-            continue;
-          }
-          if (!isInterestingAlloca(*AI))
-            continue;
-          if (II->getIntrinsicID() == Intrinsic::lifetime_start)
-            AllocasToInstrument[AI].LifetimeStart.push_back(II);
-          else
-            AllocasToInstrument[AI].LifetimeEnd.push_back(II);
+        if (!isInterestingAlloca(*AI))
           continue;
-        }
+        if (II->getIntrinsicID() == Intrinsic::lifetime_start)
+          AllocasToInstrument[AI].LifetimeStart.push_back(II);
+        else
+          AllocasToInstrument[AI].LifetimeEnd.push_back(II);
+        continue;
       }
+    }
 
-      Instruction *ExitUntag = getUntagLocationIfFunctionExit(Inst);
-      if (ExitUntag)
-        RetVec.push_back(ExitUntag);
+    Instruction *ExitUntag = getUntagLocationIfFunctionExit(Inst);
+    if (ExitUntag)
+      RetVec.push_back(ExitUntag);
 
-      if (auto *DVI = dyn_cast<DbgVariableIntrinsic>(&Inst)) {
-        for (Value *V : DVI->location_ops()) {
-          if (auto *Alloca = dyn_cast_or_null<AllocaInst>(V))
-            if (!AllocaDbgMap.count(Alloca) ||
-                AllocaDbgMap[Alloca].back() != DVI)
-              AllocaDbgMap[Alloca].push_back(DVI);
-        }
+    if (auto *DVI = dyn_cast<DbgVariableIntrinsic>(&Inst)) {
+      for (Value *V : DVI->location_ops()) {
+        if (auto *Alloca = dyn_cast_or_null<AllocaInst>(V))
+          if (!AllocaDbgMap.count(Alloca) || AllocaDbgMap[Alloca].back() != DVI)
+            AllocaDbgMap[Alloca].push_back(DVI);
       }
+    }
 
-      if (InstrumentLandingPads && isa<LandingPadInst>(Inst))
-        LandingPadVec.push_back(&Inst);
+    if (InstrumentLandingPads && isa<LandingPadInst>(Inst))
+      LandingPadVec.push_back(&Inst);
 
-      getInterestingMemoryOperands(&Inst, OperandsToInstrument);
+    getInterestingMemoryOperands(&Inst, OperandsToInstrument);
 
-      if (MemIntrinsic *MI = dyn_cast<MemIntrinsic>(&Inst))
-        if (!ignoreMemIntrinsic(MI))
-          IntrinToInstrument.push_back(MI);
-    }
+    if (MemIntrinsic *MI = dyn_cast<MemIntrinsic>(&Inst))
+      if (!ignoreMemIntrinsic(MI))
+        IntrinToInstrument.push_back(MI);
   }
 
   initializeCallbacks(*F.getParent());
@@ -1580,16 +1578,14 @@ bool HWAddressSanitizer::sanitizeFunction(
       padInterestingAllocas(AllocasToInstrument);
 
   if (!AllocaToPaddedAllocaMap.empty()) {
-    for (auto &BB : F) {
-      for (auto &Inst : BB) {
-        if (auto *DVI = dyn_cast<DbgVariableIntrinsic>(&Inst)) {
-          SmallDenseSet<Value *> LocationOps(DVI->location_ops().begin(),
-                                             DVI->location_ops().end());
-          for (Value *V : LocationOps) {
-            if (auto *AI = dyn_cast_or_null<AllocaInst>(V)) {
-              if (auto *NewAI = AllocaToPaddedAllocaMap.lookup(AI))
-                DVI->replaceVariableLocationOp(V, NewAI);
-            }
+    for (auto &Inst : instructions(F)) {
+      if (auto *DVI = dyn_cast<DbgVariableIntrinsic>(&Inst)) {
+        SmallDenseSet<Value *> LocationOps(DVI->location_ops().begin(),
+                                           DVI->location_ops().end());
+        for (Value *V : LocationOps) {
+          if (auto *AI = dyn_cast_or_null<AllocaInst>(V)) {
+            if (auto *NewAI = AllocaToPaddedAllocaMap.lookup(AI))
+              DVI->replaceVariableLocationOp(V, NewAI);
           }
         }
       }
-- 
2.7.4