nth root - ManagedDecimal

Author: andrei-marinica

framework/base/src/types/managed/wrapped/managed_decimal.rs

@@ -139,6 +139,35 @@ impl<M: ManagedTypeApi, DECIMALS: Unsigned> From<ManagedDecimal<M, ConstDecimals
     }
 }
 
+impl<M: ManagedTypeApi, D: Decimals + Clone> ManagedDecimal<M, D> {
+    /// Integer part of the k-th root, preserving the decimal scale.
+    ///
+    /// Internally pre-scales the raw data by `scaling_factor^(k-1)` so that after
+    /// taking the integer root the decimal point lands in the correct position:
+    ///
+    /// ```text
+    /// self.data = v * 10^d
+    ///   →  scaled = self.data * (10^d)^(k-1) = v * 10^(d*k)
+    ///   →  root   = floor(scaled^(1/k)) = floor(v^(1/k) * 10^d)
+    /// ```
+    ///
+    /// Returns `0` (with the same scale) when `self` is zero.
+    ///
+    /// # Panics
+    /// Panics if `k` is zero.
+    pub fn nth_root(&self, k: u32) -> Self {
+        if k == 1 {
+            return self.clone();
+        }
+
+        let sf = self.decimals.scaling_factor::<M>();
+        // Multiply by sf^(k-1) before rooting so the decimal position is preserved.
+        // For k==0, the check in BigUint::nth_root handles the error signal.
+        let scaled = &self.data * &sf.pow(k.saturating_sub(1));
+        ManagedDecimal::from_raw_units(scaled.nth_root(k), self.decimals.clone())
+    }
+}
+
 impl<M: ManagedTypeApi> ManagedVecItem for ManagedDecimal<M, NumDecimals> {
     type PAYLOAD = ManagedVecItemPayloadBuffer<U8>; // 4 bigUint + 4 usize
 

framework/base/src/types/managed/wrapped/num/big_uint.rs

@@ -396,7 +396,11 @@ impl<M: ManagedTypeApi> BigUint<M> {
             );
 
             // new_x = (k-1)*x + self/x^(k-1)
-            api.bi_mul(new_x.get_handle(), k_minus_1_big.get_handle(), x.get_handle());
+            api.bi_mul(
+                new_x.get_handle(),
+                k_minus_1_big.get_handle(),
+                x.get_handle(),
+            );
             new_x += &x_pow_k_minus_1;
 
             // new_x /= k
@@ -405,7 +409,7 @@ impl<M: ManagedTypeApi> BigUint<M> {
             if new_x >= x {
                 break;
             }
-            
+
             // Swap handles instead of cloning: zero API calls, no allocation.
             core::mem::swap(&mut x, &mut new_x);
         }

framework/scenario/tests/managed_decimal_test.rs

@@ -660,3 +660,54 @@ pub fn test_managed_decimal_div_mix_decimals_type_reverse() {
 
     assert_eq!(result, expected)
 }
+
+// d=4 ManagedDecimal helper: raw units / 10^4
+fn md4(v: u64) -> ManagedDecimal<StaticApi, NumDecimals> {
+    ManagedDecimal::from_raw_units(BigUint::from(v), 4usize)
+}
+
+fn assert_md4_nth_root(raw: u64, k: u32, expected_raw: u64) {
+    assert_eq!(
+        md4(raw).nth_root(k).into_raw_units(),
+        &BigUint::<StaticApi>::from(expected_raw)
+    );
+}
+
+#[test]
+fn test_managed_decimal_nth_root() {
+    // k=1: identity
+    assert_md4_nth_root(40000, 1, 40000);
+
+    // zero: any k≥2 root of 0.0000 is 0.0000
+    assert_md4_nth_root(0, 2, 0);
+
+    // sqrt(4.0000) = 2.0000
+    // scaled = 40000 * 10000^1 = 400_000_000, sqrt = 20000
+    assert_md4_nth_root(40000, 2, 20000);
+
+    // sqrt(9.0000) = 3.0000
+    // scaled = 90000 * 10000 = 900_000_000, sqrt = 30000
+    assert_md4_nth_root(90000, 2, 30000);
+
+    // cbrt(8.0000) = 2.0000
+    // scaled = 80000 * 10000^2 = 8_000_000_000_000, cbrt = 20000
+    assert_md4_nth_root(80000, 3, 20000);
+
+    // cbrt(27.0000) = 3.0000
+    // scaled = 270000 * 10000^2 = 27_000_000_000_000, cbrt = 30000
+    assert_md4_nth_root(270000, 3, 30000);
+
+    // sqrt(2.0000) ≈ 1.4142 (floor)
+    // scaled = 20000 * 10000 = 200_000_000, sqrt = 14142
+    assert_md4_nth_root(20000, 2, 14142);
+
+    // cbrt(2.0000) ≈ 1.2599 (floor)
+    // scaled = 20000 * 10000^2 = 2_000_000_000_000, cbrt = 12599
+    assert_md4_nth_root(20000, 3, 12599);
+}
+
+#[test]
+#[should_panic = "StaticApi signal error: cannot compute 0th root"]
+fn test_managed_decimal_nth_root_zero_k() {
+    let _ = md4(40000).nth_root(0);
+}