Added nth root function to sdk.Decimal type

types/decimal.go

@@ -2,6 +2,7 @@ package types
 
 import (
 	"encoding/json"
+	"errors"
 	"fmt"
 	"math/big"
 	"strconv"
@@ -261,7 +262,6 @@ func (d Dec) MulInt64(i int64) Dec {
 
 // quotient
 func (d Dec) Quo(d2 Dec) Dec {
-
 	// multiply precision twice
 	mul := new(big.Int).Mul(d.Int, precisionReuse)
 	mul.Mul(mul, precisionReuse)
@@ -318,30 +318,74 @@ func (d Dec) QuoInt64(i int64) Dec {
 	return Dec{mul}
 }
 
-// ApproxSqrt returns an approximate sqrt estimation using Newton's method to
-// compute square roots x=√d for d > 0. The algorithm starts with some guess and
+// ApproxRoot returns an approximate estimation of a Dec's positive real nth root
+// using Newton's method (where n is positive). The algorithm starts with some guess and
 // computes the sequence of improved guesses until an answer converges to an
-// approximate answer. It returns -(sqrt(abs(d)) if input is negative.
-func (d Dec) ApproxSqrt() Dec {
+// approximate answer.  It returns -(sqrt(abs(d)) if input is negative.
+func (d Dec) ApproxRoot(root uint64) (guess Dec, err error) {
+	defer func() {
+		if r := recover(); r != nil {
+			var ok bool
+			err, ok = r.(error)
+			if !ok {
+				err = errors.New("out of bounds")
+			}
+		}
+	}()
+
 	if d.IsNegative() {
-		return d.MulInt64(-1).ApproxSqrt().MulInt64(-1)
+		absRoot, err := d.MulInt64(-1).ApproxRoot(root)
+		return absRoot.MulInt64(-1), err
+	}
+
+	if root == 1 || d.IsZero() || d.Equal(OneDec()) {
+		return d, nil
+	}
+
+	if root == 0 {
+		return OneDec(), nil
 	}
 
-	if d.IsZero() {
-		return ZeroDec()
+	rootInt := NewIntFromUint64(root)
+	guess, delta := OneDec(), OneDec()
+
+	for delta.Abs().GT(SmallestDec()) {
+		prev := guess.Power(root - 1)
+		if prev.IsZero() {
+			prev = SmallestDec()
+		}
+		delta = d.Quo(prev)
+		delta = delta.Sub(guess)
+		delta = delta.QuoInt(rootInt)
+
+		guess = guess.Add(delta)
 	}
 
-	z := OneDec()
-	// first guess
-	z = z.Sub((z.Mul(z).Sub(d)).Quo(z.MulInt64(2)))
+	return guess, nil
+}
 
-	// iterate until change is very small
-	for zNew, delta := z, z; delta.GT(SmallestDec()); z = zNew {
-		zNew = zNew.Sub((zNew.Mul(zNew).Sub(d)).Quo(zNew.MulInt64(2)))
-		delta = z.Sub(zNew)
+// Power returns a the result of raising to a positive integer power
+func (d Dec) Power(power uint64) Dec {
+	if power == 0 {
+		return OneDec()
 	}
+	tmp := OneDec()
+	for i := power; i > 1; {
+		if i%2 == 0 {
+			i /= 2
+		} else {
+			tmp = tmp.Mul(d)
+			i = (i - 1) / 2
+		}
+		d = d.Mul(d)
+	}
+	return d.Mul(tmp)
+}
 
-	return z
+// ApproxSqrt is a wrapper around ApproxRoot for the common special case
+// of finding square root of a number.
+func (d Dec) ApproxSqrt() (Dec, error) {
+	return d.ApproxRoot(2)
 }
 
 // is integer, e.g. decimals are zero

types/decimal_test.go

@@ -1,6 +1,7 @@
 package types
 
 import (
+	"fmt"
 	"math/big"
 	"testing"
 
@@ -425,6 +426,48 @@ func TestDecCeil(t *testing.T) {
 	}
 }
 
+func TestPower(t *testing.T) {
+	testCases := []struct {
+		input    Dec
+		power    uint64
+		expected Dec
+	}{
+		{OneDec(), 10, OneDec()},                                               // 1.0 ^ (10) => 1.0
+		{NewDecWithPrec(5, 1), 2, NewDecWithPrec(25, 2)},                       // 0.5 ^ 2 => 0.25
+		{NewDecWithPrec(2, 1), 2, NewDecWithPrec(4, 2)},                        // 0.2 ^ 2 => 0.04
+		{NewDecFromInt(NewInt(3)), 3, NewDecFromInt(NewInt(27))},               // 3 ^ 3 => 27
+		{NewDecFromInt(NewInt(-3)), 4, NewDecFromInt(NewInt(81))},              // -3 ^ 4 = 81
+		{NewDecWithPrec(1414213562373095049, 18), 2, NewDecFromInt(NewInt(2))}, // 1.414213562373095049 ^ 2 = 2
+	}
+
+	for i, tc := range testCases {
+		res := tc.input.Power(tc.power)
+		require.True(t, tc.expected.Sub(res).Abs().LTE(SmallestDec()), "unexpected result for test case %d, input: %v", i, tc.input)
+	}
+}
+
+func TestApproxRoot(t *testing.T) {
+	testCases := []struct {
+		input    Dec
+		root     uint64
+		expected Dec
+	}{
+		{OneDec(), 10, OneDec()},                                               // 1.0 ^ (0.1) => 1.0
+		{NewDecWithPrec(25, 2), 2, NewDecWithPrec(5, 1)},                       // 0.25 ^ (0.5) => 0.5
+		{NewDecWithPrec(4, 2), 2, NewDecWithPrec(2, 1)},                        // 0.04 => 0.2
+		{NewDecFromInt(NewInt(27)), 3, NewDecFromInt(NewInt(3))},               // 27 ^ (1/3) => 3
+		{NewDecFromInt(NewInt(-81)), 4, NewDecFromInt(NewInt(-3))},             // -81 ^ (0.25) => -3
+		{NewDecFromInt(NewInt(2)), 2, NewDecWithPrec(1414213562373095049, 18)}, // 2 ^ (0.5) => 1.414213562373095049
+		{NewDecWithPrec(1005, 3), 31536000, MustNewDecFromStr("1.000000000158153904")},
+	}
+
+	for i, tc := range testCases {
+		res, err := tc.input.ApproxRoot(tc.root)
+		require.Nil(t, err)
+		require.True(t, tc.expected.Sub(res).Abs().LTE(SmallestDec()), "unexpected result for test case %d, input: %v", i, tc.input)
+	}
+}
+
 func TestApproxSqrt(t *testing.T) {
 	testCases := []struct {
 		input    Dec
@@ -439,7 +482,8 @@ func TestApproxSqrt(t *testing.T) {
 	}
 
 	for i, tc := range testCases {
-		res := tc.input.ApproxSqrt()
+		res, err := tc.input.ApproxSqrt()
+		require.Nil(t, err)
 		require.Equal(t, tc.expected, res, "unexpected result for test case %d, input: %v", i, tc.input)
 	}
 }

types/int.go

@@ -114,6 +114,13 @@ func NewInt(n int64) Int {
 	return Int{big.NewInt(n)}
 }
 
+// NewInt constructs Int from int64
+func NewIntFromUint64(n uint64) Int {
+	b := big.NewInt(0)
+	b.SetUint64(n)
+	return Int{b}
+}
+
 // NewIntFromBigInt constructs Int from big.Int
 func NewIntFromBigInt(i *big.Int) Int {
 	if i.BitLen() > maxBitLen {