Skip to content

Commit

Permalink
refactor(math): refactor ApproxRoot for readality (#22263)
Browse files Browse the repository at this point in the history
Co-authored-by: heren-ke <[email protected]>
  • Loading branch information
xujiangyu and heren-ke authored Oct 30, 2024
1 parent fa3eadc commit ee16adb
Show file tree
Hide file tree
Showing 2 changed files with 17 additions and 6 deletions.
20 changes: 14 additions & 6 deletions math/dec.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
3 changes: 3 additions & 0 deletions math/dec_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit ee16adb

Please sign in to comment.