1 files changed, 13 insertions, 3 deletions
diff --git a/llvm/lib/Analysis/CodeMetrics.cpp b/llvm/lib/Analysis/CodeMetrics.cpp
index 157811c04eb5..8c8e2ee6627f 100644
--- a/llvm/lib/Analysis/CodeMetrics.cpp
+++ b/llvm/lib/Analysis/CodeMetrics.cpp
@@ -18,6 +18,7 @@
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/Function.h"
 #include "llvm/Support/Debug.h"
+#include "llvm/Support/InstructionCost.h"
 
 #define DEBUG_TYPE "code-metrics"
 
@@ -116,7 +117,14 @@ void CodeMetrics::analyzeBasicBlock(
     const BasicBlock *BB, const TargetTransformInfo &TTI,
     const SmallPtrSetImpl<const Value *> &EphValues, bool PrepareForLTO) {
   ++NumBlocks;
-  unsigned NumInstsBeforeThisBB = NumInsts;
+  // Use a proxy variable for NumInsts of type InstructionCost, so that it can
+  // use InstructionCost's arithmetic properties such as saturation when this
+  // feature is added to InstructionCost.
+  // When storing the value back to NumInsts, we can assume all costs are Valid
+  // because the IR should not contain any nodes that cannot be costed. If that
+  // happens the cost-model is broken.
+  InstructionCost NumInstsProxy = NumInsts;
+  InstructionCost NumInstsBeforeThisBB = NumInsts;
   for (const Instruction &I : *BB) {
     // Skip ephemeral values.
     if (EphValues.count(&I))
@@ -175,7 +183,8 @@ void CodeMetrics::analyzeBasicBlock(
       if (InvI->cannotDuplicate())
         notDuplicatable = true;
 
-    NumInsts += TTI.getUserCost(&I, TargetTransformInfo::TCK_CodeSize);
+    NumInstsProxy += TTI.getUserCost(&I, TargetTransformInfo::TCK_CodeSize);
+    NumInsts = *NumInstsProxy.getValue();
   }
 
   if (isa<ReturnInst>(BB->getTerminator()))
@@ -195,5 +204,6 @@ void CodeMetrics::analyzeBasicBlock(
   notDuplicatable |= isa<IndirectBrInst>(BB->getTerminator());
 
   // Remember NumInsts for this BB.
-  NumBBInsts[BB] = NumInsts - NumInstsBeforeThisBB;
+  InstructionCost NumInstsThisBB = NumInstsProxy - NumInstsBeforeThisBB;
+  NumBBInsts[BB] = *NumInstsThisBB.getValue();
 }