Add DivMod() (#113)

chfast · web-flow · commit 78905e8fd24f · 2022-03-25T09:55:28.000+01:00
DivMod() sets z to the quotient x div y and m to the modulus x mod y and returns the pair (z, m) for y != 0. If y == 0, both z and m are set to 0 (OBS: differs from the big.Int). Closes #103.
diff --git a/uint256.go b/uint256.go
@@ -588,8 +588,20 @@ func (z *Int) Mod(x, y *Int) *Int {
 	}
 
 	var quot Int
-	rem := udivrem(quot[:], x[:], y)
-	return z.Set(&rem)
+	*z = udivrem(quot[:], x[:], y)
+	return z
+}
+
+// DivMod sets z to the quotient x div y and m to the modulus x mod y and returns the pair (z, m) for y != 0.
+// If y == 0, both z and m are set to 0 (OBS: differs from the big.Int)
+func (z *Int) DivMod(x, y, m *Int) (*Int, *Int) {
+	if y.IsZero() {
+		return z.Clear(), m.Clear()
+	}
+	var quot Int
+	*m = udivrem(quot[:], x[:], y)
+	*z = quot
+	return z, m
 }
 
 // SMod interprets x and y as two's complement signed integers,
diff --git a/uint256_test.go b/uint256_test.go
@@ -260,6 +260,8 @@ func TestRandomBinOp(t *testing.T) {
 	t.Run("Mod", func(t *testing.T) { testRandomOp(t, (*Int).Mod, bigMod) })
 	t.Run("SDiv", func(t *testing.T) { testRandomOp(t, (*Int).SDiv, bigSDiv) })
 	t.Run("SMod", func(t *testing.T) { testRandomOp(t, (*Int).SMod, bigSMod) })
+	t.Run("DivMod/Div", func(t *testing.T) { testRandomOp(t, divModDiv, bigDiv) })
+	t.Run("DivMod/Mod", func(t *testing.T) { testRandomOp(t, divModMod, bigMod) })
 	t.Run("udivrem/Div", func(t *testing.T) { testRandomOp(t, udivremDiv, bigDiv) })
 	t.Run("udivrem/Mod", func(t *testing.T) { testRandomOp(t, udivremMod, bigMod) })
 }
@@ -298,6 +300,19 @@ func TestRandomSquare(t *testing.T) {
 	)
 }
 
+// divModDiv wraps DivMod and returns quotient only
+func divModDiv(z, x, y *Int) *Int {
+	var m Int
+	z.DivMod(x, y, &m)
+	return z
+}
+
+// divModMod wraps DivMod and returns modulus only
+func divModMod(z, x, y *Int) *Int {
+	new(Int).DivMod(x, y, z)
+	return z
+}
+
 // udivremDiv wraps udivrem and returns quotient
 func udivremDiv(z, x, y *Int) *Int {
 	var quot Int
@@ -1243,6 +1258,8 @@ func TestBinOp(t *testing.T) {
 	t.Run("Mod", func(t *testing.T) { proc(t, (*Int).Mod, bigMod) })
 	t.Run("SDiv", func(t *testing.T) { proc(t, (*Int).SDiv, bigSDiv) })
 	t.Run("SMod", func(t *testing.T) { proc(t, (*Int).SMod, bigSMod) })
+	t.Run("DivMod/Div", func(t *testing.T) { proc(t, divModDiv, bigDiv) })
+	t.Run("DivMod/Mod", func(t *testing.T) { proc(t, divModMod, bigMod) })
 	t.Run("udivrem/Div", func(t *testing.T) { proc(t, udivremDiv, bigDiv) })
 	t.Run("udivrem/Mod", func(t *testing.T) { proc(t, udivremMod, bigMod) })
 	t.Run("Exp", func(t *testing.T) { proc(t, (*Int).Exp, bigExp) })
