1 files changed, 16 insertions, 2 deletions
diff --git a/include/llvm/Support/BranchProbability.h b/include/llvm/Support/BranchProbability.h
index 05c24d4fcf..a1e5ceb625 100644
--- a/include/llvm/Support/BranchProbability.h
+++ b/include/llvm/Support/BranchProbability.h
@@ -15,6 +15,7 @@
 #define LLVM_SUPPORT_BRANCHPROBABILITY_H
 
 #include "llvm/Support/DataTypes.h"
+#include <cassert>
 
 namespace llvm {
 
@@ -22,15 +23,21 @@ class raw_ostream;
 
 // This class represents Branch Probability as a non-negative fraction.
 class BranchProbability {
-
   // Numerator
   uint32_t N;
 
   // Denominator
   uint32_t D;
 
+  int64_t compare(BranchProbability RHS) const {
+    return (uint64_t)N * RHS.D - (uint64_t)D * RHS.N;
+  }
+
 public:
-  BranchProbability(uint32_t n, uint32_t d);
+  BranchProbability(uint32_t n, uint32_t d) : N(n), D(d) {
+    assert(d > 0 && "Denomiator cannot be 0!");
+    assert(n <= d && "Probability cannot be bigger than 1!");
+  }
 
   uint32_t getNumerator() const { return N; }
   uint32_t getDenominator() const { return D; }
@@ -43,6 +50,13 @@ public:
   void print(raw_ostream &OS) const;
 
   void dump() const;
+
+  bool operator==(BranchProbability RHS) const { return compare(RHS) == 0; }
+  bool operator!=(BranchProbability RHS) const { return compare(RHS) != 0; }
+  bool operator< (BranchProbability RHS) const { return compare(RHS) <  0; }
+  bool operator> (BranchProbability RHS) const { return compare(RHS) >  0; }
+  bool operator<=(BranchProbability RHS) const { return compare(RHS) <= 0; }
+  bool operator>=(BranchProbability RHS) const { return compare(RHS) >= 0; }
 };
 
 raw_ostream &operator<<(raw_ostream &OS, const BranchProbability &Prob);