diff --git a/math/dec.go b/math/dec.go index 3ed930d5c484..76a37113b6b8 100644 --- a/math/dec.go +++ b/math/dec.go @@ -467,33 +467,41 @@ func (d LegacyDec) ApproxRoot(root uint64) (guess LegacyDec, err error) { } }() + if root == 0 { + // Return 1 as root 0 of any number is considered 1. + return LegacyOneDec(), nil + } + if d.IsNegative() { absRoot, err := d.Neg().ApproxRoot(root) return absRoot.NegMut(), err } - // One decimal, that we invalidate later. Helps us save a heap allocation. + // Direct return for base cases: d^1 = d or when d is 0 or 1. scratchOneDec := LegacyOneDec() if root == 1 || d.IsZero() || d.Equal(scratchOneDec) { return d, nil } - if root == 0 { - return scratchOneDec, nil - } - guess, delta := scratchOneDec, LegacyOneDec() - for iter := 0; iter < maxApproxRootIterations && delta.Abs().GT(smallestDec); iter++ { + for iter := 0; iter < maxApproxRootIterations; iter++ { prev := guess.Power(root - 1) if prev.IsZero() { prev = smallestDec } + + // Compute delta = (d/prev - guess) / root delta.Set(d).QuoMut(prev) delta.SubMut(guess) delta.QuoInt64Mut(int64(root)) guess.AddMut(delta) + + // Stop when delta is small enough + if delta.Abs().LTE(smallestDec) { + break + } } return guess, nil diff --git a/math/dec_test.go b/math/dec_test.go index 90d610199f3f..316b6aaa63dc 100644 --- a/math/dec_test.go +++ b/math/dec_test.go @@ -449,6 +449,9 @@ func (s *decimalTestSuite) TestApproxRoot() { root uint64 expected math.LegacyDec }{ + {math.LegacyNewDecFromInt(math.NewInt(2)), 0, math.LegacyOneDec()}, // 2 ^ 0 => 1.0 + {math.LegacyNewDecWithPrec(4, 2), 0, math.LegacyOneDec()}, // 0.04 ^ 0 => 1.0 + {math.LegacyNewDec(0), 1, math.LegacyNewDec(0)}, // 0 ^ 1 => 0 {math.LegacyOneDec(), 10, math.LegacyOneDec()}, // 1.0 ^ (0.1) => 1.0 {math.LegacyNewDecWithPrec(25, 2), 2, math.LegacyNewDecWithPrec(5, 1)}, // 0.25 ^ (0.5) => 0.5 {math.LegacyNewDecWithPrec(4, 2), 2, math.LegacyNewDecWithPrec(2, 1)}, // 0.04 ^ (0.5) => 0.2