3 files changed, 68 insertions, 14 deletions

diff --git a/include/llvm/Analysis/BlockFrequencyInfoImpl.h b/include/llvm/Analysis/BlockFrequencyInfoImpl.h
index 7ad76e5350..264aff3725 100644
--- a/include/llvm/Analysis/BlockFrequencyInfoImpl.h
+++ b/include/llvm/Analysis/BlockFrequencyInfoImpl.h
@@ -327,20 +327,16 @@ private:
     return countLeadingZeros32(Digits) + Width - 32;
   }
 
-  static UnsignedFloat adjustToWidth(uint64_t N, int32_t S) {
-    assert(S >= MinExponent);
-    assert(S <= MaxExponent);
-    if (Width == 64 || N <= DigitsLimits::max())
-      return UnsignedFloat(N, S);
-
-    // Shift right.
-    int Shift = 64 - Width - countLeadingZeros64(N);
-    DigitsType Shifted = N >> Shift;
-
-    // Round.
-    assert(S + Shift <= MaxExponent);
-    return getRounded(UnsignedFloat(Shifted, S + Shift),
-                      N & UINT64_C(1) << (Shift - 1));
+  /// \brief Adjust a number to width, rounding up if necessary.
+  ///
+  /// Should only be called for \c Shift close to zero.
+  ///
+  /// \pre Shift >= MinExponent && Shift + 64 <= MaxExponent.
+  static UnsignedFloat adjustToWidth(uint64_t N, int32_t Shift) {
+    assert(Shift >= MinExponent && "Shift should be close to 0");
+    assert(Shift <= MaxExponent - 64 && "Shift should be close to 0");
+    auto Adjusted = ScaledNumbers::getAdjusted<DigitsT>(N, Shift);
+    return Adjusted;
   }
 
   static UnsignedFloat getRounded(UnsignedFloat P, bool Round) {
diff --git a/include/llvm/Support/ScaledNumber.h b/include/llvm/Support/ScaledNumber.h
index afdc7dc188..54b5386c5a 100644
--- a/include/llvm/Support/ScaledNumber.h
+++ b/include/llvm/Support/ScaledNumber.h
@@ -22,6 +22,8 @@
 #ifndef LLVM_SUPPORT_SCALEDNUMBER_H
 #define LLVM_SUPPORT_SCALEDNUMBER_H
 
+#include "llvm/Support/MathExtras.h"
+
 #include <cstdint>
 #include <limits>
 #include <utility>
@@ -50,6 +52,39 @@ inline std::pair<DigitsT, int16_t> getRounded(DigitsT Digits, int16_t Scale,
       return std::make_pair(DigitsT(1) << (getWidth<DigitsT>() - 1), Scale + 1);
   return std::make_pair(Digits, Scale);
 }
+
+/// \brief Adjust a 64-bit scaled number down to the appropriate width.
+///
+/// Adjust a soft float with 64-bits of digits down, keeping as much
+/// information as possible, and rounding up on half.
+///
+/// \pre Adding 1 to \c Scale will not overflow INT16_MAX.
+template <class DigitsT>
+inline std::pair<DigitsT, int16_t> getAdjusted(uint64_t Digits,
+                                               int16_t Scale = 0) {
+  static_assert(!std::numeric_limits<DigitsT>::is_signed, "expected unsigned");
+
+  const int Width = getWidth<DigitsT>();
+  if (Width == 64 || Digits <= std::numeric_limits<DigitsT>::max())
+    return std::make_pair(Digits, Scale);
+
+  // Shift right and round.
+  int Shift = 64 - Width - countLeadingZeros(Digits);
+  return getRounded<DigitsT>(Digits >> Shift, Scale + Shift,
+                             Digits & (UINT64_C(1) << (Shift - 1)));
+}
+
+/// \brief Convenience helper for adjusting to 32 bits.
+inline std::pair<uint32_t, int16_t> getAdjusted32(uint64_t Digits,
+                                                  int16_t Scale = 0) {
+  return getAdjusted<uint32_t>(Digits, Scale);
+}
+
+/// \brief Convenience helper for adjusting to 64 bits.
+inline std::pair<uint64_t, int16_t> getAdjusted64(uint64_t Digits,
+                                                  int16_t Scale = 0) {
+  return getAdjusted<uint64_t>(Digits, Scale);
+}
 }
 }
 
diff --git a/unittests/Support/ScaledNumberTest.cpp b/unittests/Support/ScaledNumberTest.cpp
index 1fa54ff81d..dac24051a6 100644
--- a/unittests/Support/ScaledNumberTest.cpp
+++ b/unittests/Support/ScaledNumberTest.cpp
@@ -57,4 +57,27 @@ TEST(ScaledNumberHelpersTest, getRounded) {
   EXPECT_EQ(getRounded<uint64_t>(UINT64_MAX, 0, true),
             SP64(UINT64_C(1) << 63, 1));
 }
+
+TEST(FloatsTest, getAdjusted) {
+  const uint64_t Max32In64 = UINT32_MAX;
+  EXPECT_EQ(getAdjusted32(0), SP32(0, 0));
+  EXPECT_EQ(getAdjusted32(0, 5), SP32(0, 5));
+  EXPECT_EQ(getAdjusted32(UINT32_MAX), SP32(UINT32_MAX, 0));
+  EXPECT_EQ(getAdjusted32(Max32In64 << 1), SP32(UINT32_MAX, 1));
+  EXPECT_EQ(getAdjusted32(Max32In64 << 1, 1), SP32(UINT32_MAX, 2));
+  EXPECT_EQ(getAdjusted32(Max32In64 << 31), SP32(UINT32_MAX, 31));
+  EXPECT_EQ(getAdjusted32(Max32In64 << 32), SP32(UINT32_MAX, 32));
+  EXPECT_EQ(getAdjusted32(Max32In64 + 1), SP32(1u << 31, 1));
+  EXPECT_EQ(getAdjusted32(UINT64_MAX), SP32(1u << 31, 33));
+
+  EXPECT_EQ(getAdjusted64(0), SP64(0, 0));
+  EXPECT_EQ(getAdjusted64(0, 5), SP64(0, 5));
+  EXPECT_EQ(getAdjusted64(UINT32_MAX), SP64(UINT32_MAX, 0));
+  EXPECT_EQ(getAdjusted64(Max32In64 << 1), SP64(Max32In64 << 1, 0));
+  EXPECT_EQ(getAdjusted64(Max32In64 << 1, 1), SP64(Max32In64 << 1, 1));
+  EXPECT_EQ(getAdjusted64(Max32In64 << 31), SP64(Max32In64 << 31, 0));
+  EXPECT_EQ(getAdjusted64(Max32In64 << 32), SP64(Max32In64 << 32, 0));
+  EXPECT_EQ(getAdjusted64(Max32In64 + 1), SP64(Max32In64 + 1, 0));
+  EXPECT_EQ(getAdjusted64(UINT64_MAX), SP64(UINT64_MAX, 0));
+}
 }