Add constexpr_add/sub_overflow_clamped in folly/ConstexprMath.h

Summary: Provide UBSAN-clean add/sub function on integral types. Helps to avoid overflow.

Reviewed By: yfeldblum

Differential Revision: D7230723

fbshipit-source-id: 10fa54b405f9643dfbbb1940cf100516fd1341e1

folly/ConstexprMath.h

@@ -16,6 +16,8 @@
 
 #pragma once
 
+#include <cstdint>
+#include <limits>
 #include <type_traits>
 
 namespace folly {
@@ -146,4 +148,55 @@ constexpr T constexpr_pow(T base, std::size_t exp) {
               (exp % 2 ? base : T(1));
 }
 
+template <typename T>
+constexpr T constexpr_add_overflow_clamped(T a, T b) {
+  using L = std::numeric_limits<T>;
+  using M = std::intmax_t;
+  static_assert(
+      !std::is_integral<T>::value || sizeof(T) <= sizeof(M),
+      "Integral type too large!");
+  // clang-format off
+  return
+    // don't do anything special for non-integral types.
+    !std::is_integral<T>::value ? a + b :
+    // for narrow integral types, just convert to intmax_t.
+    sizeof(T) < sizeof(M)
+      ? T(constexpr_clamp(M(a) + M(b), M(L::min()), M(L::max()))) :
+    // when a >= 0, cannot add more than `MAX - a` onto a.
+    !(a < 0) ? a + constexpr_min(b, T(L::max() - a)) :
+    // a < 0 && b >= 0, `a + b` will always be in valid range of type T.
+    !(b < 0) ? a + b :
+    // a < 0 && b < 0, keep the result >= MIN.
+               a + constexpr_max(b, T(L::min() - a));
+  // clang-format on
+}
+
+template <typename T>
+constexpr T constexpr_sub_overflow_clamped(T a, T b) {
+  using L = std::numeric_limits<T>;
+  using M = std::intmax_t;
+  static_assert(
+      !std::is_integral<T>::value || sizeof(T) <= sizeof(M),
+      "Integral type too large!");
+  // clang-format off
+  return
+    // don't do anything special for non-integral types.
+    !std::is_integral<T>::value ? a - b :
+    // for unsigned type, keep result >= 0.
+    std::is_unsigned<T>::value ? (a < b ? 0 : a - b) :
+    // for narrow signed integral types, just convert to intmax_t.
+    sizeof(T) < sizeof(M)
+      ? T(constexpr_clamp(M(a) - M(b), M(L::min()), M(L::max()))) :
+    // (a >= 0 && b >= 0) || (a < 0 && b < 0), `a - b` will always be valid.
+    (a < 0) == (b < 0) ? a - b :
+    // MIN < b, so `-b` should be in valid range (-MAX <= -b <= MAX),
+    // convert subtraction to addition.
+    L::min() < b ? constexpr_add_overflow_clamped(a, T(-b)) :
+    // -b = -MIN = (MAX + 1) and a <= -1, result is in valid range.
+    a < 0 ? a - b :
+    // -b = -MIN = (MAX + 1) and a >= 0, result > MAX.
+            L::max();
+  // clang-format on
+}
+
 } // namespace folly

folly/test/ConstexprMathBenchmark.cpp

+/*
+ * Copyright 2018-present Facebook, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <folly/Benchmark.h>
+#include <folly/ConstexprMath.h>
+#include <glog/logging.h>
+#include <limits>
+#include <type_traits>
+
+template <typename ValueT>
+constexpr ValueT UBSafeAdd(ValueT a, ValueT b) {
+  using UnsignedT = typename std::make_unsigned<ValueT>::type;
+  return static_cast<ValueT>(
+      static_cast<UnsignedT>(a) + static_cast<UnsignedT>(b));
+}
+
+template <typename ValueT>
+constexpr ValueT UBSafeSub(ValueT a, ValueT b) {
+  using UnsignedT = typename std::make_unsigned<ValueT>::type;
+  return static_cast<ValueT>(
+      static_cast<UnsignedT>(a) - static_cast<UnsignedT>(b));
+}
+
+template <typename ValueT, typename Op>
+void Run(size_t iterations, ValueT kMin, ValueT kMax, Op&& op) {
+  auto kMid = (kMin + kMax) / 2;
+
+  for (size_t round = 0; round < iterations; round++) {
+    for (ValueT a = kMin; a < kMin + 100; a++) {
+      for (ValueT b = kMin; b < kMin + 100; b++) {
+        auto a1 = a, b1 = b;
+        folly::makeUnpredictable(a1);
+        folly::makeUnpredictable(b1);
+        ValueT c = op(a1, b1);
+        folly::doNotOptimizeAway(c);
+      }
+    }
+    for (ValueT a = kMin; a < kMin + 100; a++) {
+      for (ValueT b = kMid - 50; b < kMid + 50; b++) {
+        auto a1 = a, b1 = b;
+        folly::makeUnpredictable(a1);
+        folly::makeUnpredictable(b1);
+        ValueT c = op(a1, b1);
+        folly::doNotOptimizeAway(c);
+      }
+    }
+    for (ValueT a = kMin; a < kMin + 100; a++) {
+      for (ValueT b = kMax - 100; b < kMax; b++) {
+        auto a1 = a, b1 = b;
+        folly::makeUnpredictable(a1);
+        folly::makeUnpredictable(b1);
+        ValueT c = op(a1, b1);
+        folly::doNotOptimizeAway(c);
+      }
+    }
+    for (ValueT a = kMid - 50; a < kMid + 50; a++) {
+      for (ValueT b = kMin; b < kMin + 100; b++) {
+        auto a1 = a, b1 = b;
+        folly::makeUnpredictable(a1);
+        folly::makeUnpredictable(b1);
+        ValueT c = op(a1, b1);
+        folly::doNotOptimizeAway(c);
+      }
+    }
+    for (ValueT a = kMid - 50; a < kMid + 50; a++) {
+      for (ValueT b = kMid - 50; b < kMid + 50; b++) {
+        auto a1 = a, b1 = b;
+        folly::makeUnpredictable(a1);
+        folly::makeUnpredictable(b1);
+        ValueT c = op(a1, b1);
+        folly::doNotOptimizeAway(c);
+      }
+    }
+    for (ValueT a = kMid - 50; a < kMid + 50; a++) {
+      for (ValueT b = kMax - 100; b < kMax; b++) {
+        auto a1 = a, b1 = b;
+        folly::makeUnpredictable(a1);
+        folly::makeUnpredictable(b1);
+        ValueT c = op(a1, b1);
+        folly::doNotOptimizeAway(c);
+      }
+    }
+    for (ValueT a = kMax - 100; a < kMax; a++) {
+      for (ValueT b = kMin; b < kMin + 100; b++) {
+        auto a1 = a, b1 = b;
+        folly::makeUnpredictable(a1);
+        folly::makeUnpredictable(b1);
+        ValueT c = op(a1, b1);
+        folly::doNotOptimizeAway(c);
+      }
+    }
+    for (ValueT a = kMax - 100; a < kMax; a++) {
+      for (ValueT b = kMid - 50; b < kMid + 50; b++) {
+        auto a1 = a, b1 = b;
+        folly::makeUnpredictable(a1);
+        folly::makeUnpredictable(b1);
+        ValueT c = op(a1, b1);
+        folly::doNotOptimizeAway(c);
+      }
+    }
+    for (ValueT a = kMax - 100; a < kMax; a++) {
+      for (ValueT b = kMax - 100; b < kMax; b++) {
+        auto a1 = a, b1 = b;
+        folly::makeUnpredictable(a1);
+        folly::makeUnpredictable(b1);
+        ValueT c = op(a1, b1);
+        folly::doNotOptimizeAway(c);
+      }
+    }
+  }
+}
+
+template <typename ValueT>
+void Add(size_t iterations, ValueT kMin, ValueT kMax) {
+  Run<ValueT>(iterations, kMin, kMax, [](ValueT a, ValueT b) {
+    return UBSafeAdd(a, b);
+  });
+}
+
+template <typename ValueT>
+void NoOverflowAdd(size_t iterations, ValueT kMin, ValueT kMax) {
+  Run<ValueT>(iterations, kMin, kMax, [](ValueT a, ValueT b) {
+    return folly::constexpr_add_overflow_clamped(a, b);
+  });
+}
+
+template <typename ValueT>
+void Sub(size_t iterations, ValueT kMin, ValueT kMax) {
+  Run<ValueT>(iterations, kMin, kMax, [](ValueT a, ValueT b) {
+    return UBSafeSub(a, b);
+  });
+}
+
+template <typename ValueT>
+void NoOverflowSub(size_t iterations, ValueT kMin, ValueT kMax) {
+  Run<ValueT>(iterations, kMin, kMax, [](ValueT a, ValueT b) {
+    return folly::constexpr_sub_overflow_clamped(a, b);
+  });
+}
+
+#define GENERATE_BENCHMARKS_FOR_TYPE(ValueT) \
+  BENCHMARK_NAMED_PARAM(                     \
+      Add,                                   \
+      ValueT,                                \
+      std::numeric_limits<ValueT>::min(),    \
+      std::numeric_limits<ValueT>::max());   \
+  BENCHMARK_RELATIVE_NAMED_PARAM(            \
+      NoOverflowAdd,                         \
+      ValueT,                                \
+      std::numeric_limits<ValueT>::min(),    \
+      std::numeric_limits<ValueT>::max());   \
+  BENCHMARK_NAMED_PARAM(                     \
+      Sub,                                   \
+      ValueT,                                \
+      std::numeric_limits<ValueT>::min(),    \
+      std::numeric_limits<ValueT>::max());   \
+  BENCHMARK_RELATIVE_NAMED_PARAM(            \
+      NoOverflowSub,                         \
+      ValueT,                                \
+      std::numeric_limits<ValueT>::min(),    \
+      std::numeric_limits<ValueT>::max())
+
+GENERATE_BENCHMARKS_FOR_TYPE(int8_t);
+BENCHMARK_DRAW_LINE()
+GENERATE_BENCHMARKS_FOR_TYPE(uint8_t);
+BENCHMARK_DRAW_LINE()
+GENERATE_BENCHMARKS_FOR_TYPE(int16_t);
+BENCHMARK_DRAW_LINE()
+GENERATE_BENCHMARKS_FOR_TYPE(uint16_t);
+BENCHMARK_DRAW_LINE()
+GENERATE_BENCHMARKS_FOR_TYPE(int32_t);
+BENCHMARK_DRAW_LINE()
+GENERATE_BENCHMARKS_FOR_TYPE(uint32_t);
+BENCHMARK_DRAW_LINE()
+GENERATE_BENCHMARKS_FOR_TYPE(int64_t);
+BENCHMARK_DRAW_LINE()
+GENERATE_BENCHMARKS_FOR_TYPE(uint64_t);
+
+int main(int argc, char** argv) {
+  gflags::ParseCommandLineFlags(&argc, &argv, true);
+  folly::runBenchmarks();
+  return 0;
+}

folly/test/ConstexprMathTest.cpp

@@ -193,3 +193,110 @@ TEST_F(ConstexprMathTest, constexpr_pow) {
     EXPECT_EQ(64, a);
   }
 }
+
+constexpr auto kInt64Max = std::numeric_limits<int64_t>::max();
+constexpr auto kInt64Min = std::numeric_limits<int64_t>::min();
+constexpr auto kUInt64Max = std::numeric_limits<uint64_t>::max();
+constexpr auto kInt8Max = std::numeric_limits<int8_t>::max();
+constexpr auto kInt8Min = std::numeric_limits<int8_t>::min();
+constexpr auto kUInt8Max = std::numeric_limits<uint8_t>::max();
+
+TEST_F(ConstexprMathTest, constexpr_add_overflow_clamped) {
+  for (int a = kInt8Min; a <= kInt8Max; a++) {
+    for (int b = kInt8Min; b <= kInt8Max; b++) {
+      int c = folly::constexpr_clamp(a + b, int(kInt8Min), int(kInt8Max));
+      int8_t a1 = a;
+      int8_t b1 = b;
+      int8_t c1 = folly::constexpr_add_overflow_clamped(a1, b1);
+      ASSERT_LE(c1, kInt8Max);
+      ASSERT_GE(c1, kInt8Min);
+      ASSERT_EQ(c1, c);
+    }
+  }
+
+  for (int a = 0; a <= kUInt8Max; a++) {
+    for (int b = 0; b <= kUInt8Max; b++) {
+      int c = folly::constexpr_clamp(a + b, 0, int(kUInt8Max));
+      uint8_t a1 = a;
+      uint8_t b1 = b;
+      uint8_t c1 = folly::constexpr_add_overflow_clamped(a1, b1);
+      ASSERT_LE(c1, kUInt8Max);
+      ASSERT_GE(c1, 0);
+      ASSERT_EQ(c1, c);
+    }
+  }
+
+  constexpr auto v1 =
+      folly::constexpr_add_overflow_clamped(int64_t(23), kInt64Max - 12);
+  EXPECT_EQ(kInt64Max, v1);
+
+  constexpr auto v2 =
+      folly::constexpr_add_overflow_clamped(int64_t(23), int64_t(12));
+  EXPECT_EQ(int64_t(35), v2);
+
+  constexpr auto v3 =
+      folly::constexpr_add_overflow_clamped(int64_t(-23), int64_t(12));
+  EXPECT_EQ(int64_t(-11), v3);
+
+  constexpr auto v4 =
+      folly::constexpr_add_overflow_clamped(int64_t(-23), int64_t(-12));
+  EXPECT_EQ(int64_t(-35), v4);
+
+  constexpr auto v5 =
+      folly::constexpr_add_overflow_clamped(uint64_t(23), kUInt64Max - 12);
+  EXPECT_EQ(kUInt64Max, v5);
+}
+
+TEST_F(ConstexprMathTest, constexpr_sub_overflow_clamped) {
+  for (int a = kInt8Min; a <= kInt8Max; a++) {
+    for (int b = kInt8Min; b <= kInt8Max; b++) {
+      int c = folly::constexpr_clamp(a - b, int(kInt8Min), int(kInt8Max));
+      int8_t a1 = a;
+      int8_t b1 = b;
+      int8_t c1 = folly::constexpr_sub_overflow_clamped(a1, b1);
+      ASSERT_LE(c1, kInt8Max);
+      ASSERT_GE(c1, kInt8Min);
+      ASSERT_EQ(c1, c);
+    }
+  }
+
+  for (int a = 0; a <= kUInt8Max; a++) {
+    for (int b = 0; b <= kUInt8Max; b++) {
+      int c = folly::constexpr_clamp(a - b, 0, int(kUInt8Max));
+      uint8_t a1 = a;
+      uint8_t b1 = b;
+      uint8_t c1 = folly::constexpr_sub_overflow_clamped(a1, b1);
+      ASSERT_LE(c1, kUInt8Max);
+      ASSERT_GE(c1, 0);
+      ASSERT_EQ(c1, c);
+    }
+  }
+
+  constexpr auto v1 =
+      folly::constexpr_sub_overflow_clamped(int64_t(23), int64_t(12));
+  EXPECT_EQ(int64_t(11), v1);
+
+  constexpr auto v2 =
+      folly::constexpr_sub_overflow_clamped(int64_t(-23), int64_t(-12));
+  EXPECT_EQ(int64_t(-11), v2);
+
+  constexpr auto v3 =
+      folly::constexpr_sub_overflow_clamped(int64_t(23), int64_t(-12));
+  EXPECT_EQ(int64_t(35), v3);
+
+  constexpr auto v4 =
+      folly::constexpr_sub_overflow_clamped(int64_t(23), kInt64Min);
+  EXPECT_EQ(kInt64Max, v4);
+
+  constexpr auto v5 =
+      folly::constexpr_sub_overflow_clamped(int64_t(-23), kInt64Min);
+  EXPECT_EQ(kInt64Max - 22, v5);
+
+  constexpr auto v6 =
+      folly::constexpr_sub_overflow_clamped(uint64_t(23), uint64_t(12));
+  EXPECT_EQ(uint64_t(11), v6);
+
+  constexpr auto v7 =
+      folly::constexpr_sub_overflow_clamped(uint64_t(12), uint64_t(23));
+  EXPECT_EQ(uint64_t(0), v7);
+}