From 374651d9aaa9eab5022993f2685af4f069a0f253 Mon Sep 17 00:00:00 2001
From: Michael Zolotukhin <mzolotukhin@apple.com>
Date: Fri, 26 Feb 2016 01:44:04 +0000
Subject: [PATCH] [UnitTests] UnrollAnalyzer: make unit-test more general so
 that it can cover more cases in future.

llvm-svn: 261954
---
 llvm/unittests/Analysis/UnrollAnalyzer.cpp | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/llvm/unittests/Analysis/UnrollAnalyzer.cpp b/llvm/unittests/Analysis/UnrollAnalyzer.cpp
index 79800ef10868..464bbdfe28f9 100644
--- a/llvm/unittests/Analysis/UnrollAnalyzer.cpp
+++ b/llvm/unittests/Analysis/UnrollAnalyzer.cpp
@@ -32,14 +32,16 @@ struct UnrollAnalyzerTest : public FunctionPass {
     FI++; // First basic block is entry - skip it.
     BasicBlock *Header = &*FI++;
     Loop *L = LI->getLoopFor(Header);
+    BasicBlock *Exiting = L->getExitingBlock();
 
     SimplifiedValuesVector.clear();
-    TripCount = SE->getSmallConstantTripCount(L, Header);
+    TripCount = SE->getSmallConstantTripCount(L, Exiting);
     for (unsigned Iteration = 0; Iteration < TripCount; Iteration++) {
       DenseMap<Value *, Constant *> SimplifiedValues;
       UnrolledInstAnalyzer Analyzer(Iteration, SimplifiedValues, *SE);
-      for (Instruction &I : *Header)
-        Analyzer.visit(I);
+      for (auto *BB : L->getBlocks())
+        for (Instruction &I : *BB)
+          Analyzer.visit(I);
       SimplifiedValuesVector.push_back(SimplifiedValues);
     }
     return false;