From ac48269f9825ddca694f4f46679734bb729198d7 Mon Sep 17 00:00:00 2001
From: Dev Ojha <ValarDragon@users.noreply.github.com>
Date: Fri, 9 Feb 2024 02:51:11 -0600
Subject: [PATCH] perf: Speedup sdk.Int overflow checks (#19386)

---
 math/CHANGELOG.md |  1 +
 math/int.go       | 37 ++++++++++++++++++++++++++++---------
 math/int_test.go  | 22 ++++++++++++++++++++++
 3 files changed, 51 insertions(+), 9 deletions(-)

diff --git a/math/CHANGELOG.md b/math/CHANGELOG.md
index d246eda970a8..85b04aed1307 100644
--- a/math/CHANGELOG.md
+++ b/math/CHANGELOG.md
@@ -42,6 +42,7 @@ Ref: https://github.com/commitizen/conventional-commit-types/blob/v3.0.0/index.j
 * [#18421](https://github.com/cosmos/cosmos-sdk/pull/18421) Add mutative api for `LegacyDec.BigInt()`.
 
 * [#18874](https://github.com/cosmos/cosmos-sdk/pull/18874) Speedup `math.Int.Mul` by removing a duplicate overflow check
+* [#19386](https://github.com/cosmos/cosmos-sdk/pull/19386) Speedup `math.Int` overflow checks
 
 ### Bug Fixes
 
diff --git a/math/int.go b/math/int.go
index 5b4bb7fa8d29..bfaf24fa3755 100644
--- a/math/int.go
+++ b/math/int.go
@@ -6,6 +6,7 @@ import (
 	"errors"
 	"fmt"
 	"math/big"
+	"math/bits"
 	"strings"
 	"sync"
 	"testing"
@@ -14,6 +15,12 @@ import (
 // MaxBitLen defines the maximum bit length supported bit Int and Uint types.
 const MaxBitLen = 256
 
+// maxWordLen defines the maximum word length supported by Int and Uint types.
+// We check overflow, by first doing a fast check if the word length is below maxWordLen
+// and if not then do the slower full bitlen check.
+// NOTE: If MaxBitLen is not a multiple of bits.UintSize, then we need to edit the used logic slightly.
+const maxWordLen = MaxBitLen / bits.UintSize
+
 // Integer errors
 var (
 	// ErrIntOverflow is the error returned when an integer overflow occurs
@@ -71,7 +78,7 @@ func unmarshalText(i *big.Int, text string) error {
 		return err
 	}
 
-	if i.BitLen() > MaxBitLen {
+	if bigIntOverflows(i) {
 		return fmt.Errorf("integer out of range: %s", text)
 	}
 
@@ -128,7 +135,7 @@ func NewIntFromBigInt(i *big.Int) Int {
 		return Int{}
 	}
 
-	if i.BitLen() > MaxBitLen {
+	if bigIntOverflows(i) {
 		panic("NewIntFromBigInt() out of bound")
 	}
 
@@ -143,7 +150,7 @@ func NewIntFromBigIntMut(i *big.Int) Int {
 		return Int{}
 	}
 
-	if i.BitLen() > MaxBitLen {
+	if bigIntOverflows(i) {
 		panic("NewIntFromBigInt() out of bound")
 	}
 
@@ -157,7 +164,7 @@ func NewIntFromString(s string) (res Int, ok bool) {
 		return
 	}
 	// Check overflow
-	if i.BitLen() > MaxBitLen {
+	if bigIntOverflows(i) {
 		ok = false
 		return
 	}
@@ -175,7 +182,7 @@ func NewIntWithDecimal(n int64, dec int) Int {
 	i.Mul(big.NewInt(n), exp)
 
 	// Check overflow
-	if i.BitLen() > MaxBitLen {
+	if bigIntOverflows(i) {
 		panic("NewIntWithDecimal() out of bound")
 	}
 	return Int{i}
@@ -285,7 +292,7 @@ func (i Int) AddRaw(i2 int64) Int {
 func (i Int) SafeAdd(i2 Int) (res Int, err error) {
 	res = Int{add(i.i, i2.i)}
 	// Check overflow
-	if res.i.BitLen() > MaxBitLen {
+	if bigIntOverflows(res.i) {
 		return Int{}, ErrIntOverflow
 	}
 	return res, nil
@@ -310,7 +317,7 @@ func (i Int) SubRaw(i2 int64) Int {
 func (i Int) SafeSub(i2 Int) (res Int, err error) {
 	res = Int{sub(i.i, i2.i)}
 	// Check overflow/underflow
-	if res.i.BitLen() > MaxBitLen {
+	if bigIntOverflows(res.i) {
 		return Int{}, ErrIntOverflow
 	}
 	return res, nil
@@ -335,7 +342,7 @@ func (i Int) MulRaw(i2 int64) Int {
 func (i Int) SafeMul(i2 Int) (res Int, err error) {
 	res = Int{mul(i.i, i2.i)}
 	// Check overflow
-	if res.i.BitLen() > MaxBitLen {
+	if bigIntOverflows(res.i) {
 		return Int{}, ErrIntOverflow
 	}
 	return res, nil
@@ -497,7 +504,7 @@ func (i *Int) Unmarshal(data []byte) error {
 		return err
 	}
 
-	if i.i.BitLen() > MaxBitLen {
+	if bigIntOverflows(i.i) {
 		return fmt.Errorf("integer out of range; got: %d, max: %d", i.i.BitLen(), MaxBitLen)
 	}
 
@@ -591,3 +598,15 @@ func FormatInt(v string) (string, error) {
 
 	return sign + sb.String(), nil
 }
+
+// check if the big int overflows.
+func bigIntOverflows(i *big.Int) bool {
+	// overflow is defined as i.BitLen() > MaxBitLen
+	// however this check can be expensive when doing many operations.
+	// So we first check if the word length is greater than maxWordLen.
+	// However the most significant word could be zero, hence we still do the bitlen check.
+	if len(i.Bits()) > maxWordLen {
+		return i.BitLen() > MaxBitLen
+	}
+	return false
+}
diff --git a/math/int_test.go b/math/int_test.go
index 1895f4bc1d3c..197f75ab6104 100644
--- a/math/int_test.go
+++ b/math/int_test.go
@@ -687,3 +687,25 @@ func BenchmarkIntSize(b *testing.B) {
 	}
 	sink = nil
 }
+
+func BenchmarkIntOverflowCheckTime(b *testing.B) {
+	var ints = []*big.Int{}
+
+	for _, st := range sizeTests {
+		ii, _ := math.NewIntFromString(st.s)
+		ints = append(ints, ii.BigInt())
+	}
+	b.ResetTimer()
+
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		for j, _ := range sizeTests {
+			got := math.NewIntFromBigIntMut(ints[j])
+			sink = got
+		}
+	}
+	if sink == nil {
+		b.Fatal("Benchmark did not run!")
+	}
+	sink = nil
+}