From ee16adb09295fdc0b98c7bfdc60d0f74a2b6b28c Mon Sep 17 00:00:00 2001
From: xujiangyu <47973791+xujiangyu@users.noreply.github.com>
Date: Wed, 30 Oct 2024 18:49:00 +0800
Subject: [PATCH] refactor(math): refactor ApproxRoot for readality (#22263)

Co-authored-by: heren-ke <fghjtuy@163.com>
---
 math/dec.go      | 20 ++++++++++++++------
 math/dec_test.go |  3 +++
 2 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/math/dec.go b/math/dec.go
index 3ed930d5c4..76a37113b6 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 90d610199f..316b6aaa63 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