From 8dfcfdef845ba4792f12b6e872c3f749e32bead4 Mon Sep 17 00:00:00 2001
From: Aaron Chen <aaronchen.lisp@gmail.com>
Date: Mon, 27 May 2024 15:58:05 +0800
Subject: [PATCH] uint256: optimize mod, DivMod (#173)

---
 benchmarks_test.go | 30 ++++++++++++++++++++++++++++++
 uint256.go         | 21 ++++++---------------
 2 files changed, 36 insertions(+), 15 deletions(-)

diff --git a/benchmarks_test.go b/benchmarks_test.go
index ab10a0c..210dfaa 100644
--- a/benchmarks_test.go
+++ b/benchmarks_test.go
@@ -705,6 +705,36 @@ func BenchmarkMod(b *testing.B) {
 	b.Run("mod256/big", func(b *testing.B) { benchmarkModBig(b, &big256Samples, &big256SamplesLt) })
 }
 
+func BenchmarkDivMod(b *testing.B) {
+	benchmarkDivModUint256 := func(b *testing.B, xSamples, modSamples *[numSamples]Int) {
+		var sink, mod Int
+		for j := 0; j < b.N; j += numSamples {
+			for i := 0; i < numSamples; i++ {
+				sink.DivMod(&xSamples[i], &modSamples[i], &mod)
+			}
+		}
+	}
+	benchmarkDivModBig := func(b *testing.B, xSamples, modSamples *[numSamples]big.Int) {
+		var sink, mod big.Int
+		for j := 0; j < b.N; j += numSamples {
+			for i := 0; i < numSamples; i++ {
+				sink.DivMod(&xSamples[i], &modSamples[i], &mod)
+			}
+		}
+	}
+
+	b.Run("small/uint256", func(b *testing.B) { benchmarkDivModUint256(b, &int32Samples, &int32SamplesLt) })
+	b.Run("mod64/uint256", func(b *testing.B) { benchmarkDivModUint256(b, &int256Samples, &int64Samples) })
+	b.Run("mod128/uint256", func(b *testing.B) { benchmarkDivModUint256(b, &int256Samples, &int128Samples) })
+	b.Run("mod192/uint256", func(b *testing.B) { benchmarkDivModUint256(b, &int256Samples, &int192Samples) })
+	b.Run("mod256/uint256", func(b *testing.B) { benchmarkDivModUint256(b, &int256Samples, &int256SamplesLt) })
+	b.Run("small/big", func(b *testing.B) { benchmarkDivModBig(b, &big32Samples, &big32SamplesLt) })
+	b.Run("mod64/big", func(b *testing.B) { benchmarkDivModBig(b, &big256Samples, &big64Samples) })
+	b.Run("mod128/big", func(b *testing.B) { benchmarkDivModBig(b, &big256Samples, &big128Samples) })
+	b.Run("mod192/big", func(b *testing.B) { benchmarkDivModBig(b, &big256Samples, &big192Samples) })
+	b.Run("mod256/big", func(b *testing.B) { benchmarkDivModBig(b, &big256Samples, &big256SamplesLt) })
+}
+
 func BenchmarkAddMod(b *testing.B) {
 	benchmarkAddModUint256 := func(b *testing.B, factorsSamples, modSamples *[numSamples]Int) {
 		iter := (b.N + numSamples - 1) / numSamples
diff --git a/uint256.go b/uint256.go
index 904bf4d..5d1df27 100644
--- a/uint256.go
+++ b/uint256.go
@@ -586,18 +586,12 @@ func (z *Int) Div(x, y *Int) *Int {
 // Mod sets z to the modulus x%y for y != 0 and returns z.
 // If y == 0, z is set to 0 (OBS: differs from the big.Int)
 func (z *Int) Mod(x, y *Int) *Int {
-	if x.IsZero() || y.IsZero() {
+	if y.IsZero() || x.Eq(y) {
 		return z.Clear()
 	}
-	switch x.Cmp(y) {
-	case -1:
-		// x < y
+	if x.Lt(y) {
 		return z.Set(x)
-	case 0:
-		// x == y
-		return z.Clear() // They are equal
 	}
-
 	// At this point:
 	// x != 0
 	// y != 0
@@ -619,15 +613,12 @@ func (z *Int) DivMod(x, y, m *Int) (*Int, *Int) {
 	if y.IsZero() {
 		return z.Clear(), m.Clear()
 	}
-
-	switch x.Cmp(y) {
-	case -1:
-		// x < y
-		return z.Clear(), m.Set(x)
-	case 0:
-		// x == y
+	if x.Eq(y) {
 		return z.SetOne(), m.Clear()
 	}
+	if x.Lt(y) {
+		return z.Clear(), m.Set(x)
+	}
 
 	// At this point:
 	// x != 0