From 0aa7cca8a1acb1f35cff32420614f7f58e38de5e Mon Sep 17 00:00:00 2001
From: water <161927344+suiwater@users.noreply.github.com>
Date: Mon, 24 Jun 2024 16:47:39 +0800
Subject: [PATCH] feat: update aliquot_sum (#309)

optimize code and add test two
---
 packages/math/src/aliquot_sum.cairo            | 11 ++++-------
 packages/math/src/tests/aliquot_sum_test.cairo |  7 +++++++
 2 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/packages/math/src/aliquot_sum.cairo b/packages/math/src/aliquot_sum.cairo
index d7bdd31d..71d8ba78 100644
--- a/packages/math/src/aliquot_sum.cairo
+++ b/packages/math/src/aliquot_sum.cairo
@@ -6,17 +6,14 @@
 /// # Returns
 /// * `felt252` - The aliquot sum of the input number.
 pub fn aliquot_sum(number: u128) -> u128 {
-    if number == 0 {
-        return 0;
-    }
-    if number == 1 {
+    if number == 0 || number == 1 {
         return 0;
     }
 
     let limit = (number / 2) + 1;
-    let mut index = 1;
-    let mut res = 0;
-    while (index != limit) {
+    let mut index = 2;
+    let mut res = 1;
+    while (index < limit) {
         if number % index == 0 {
             res = res + index;
         }
diff --git a/packages/math/src/tests/aliquot_sum_test.cairo b/packages/math/src/tests/aliquot_sum_test.cairo
index 8480607a..ce80d3d0 100644
--- a/packages/math/src/tests/aliquot_sum_test.cairo
+++ b/packages/math/src/tests/aliquot_sum_test.cairo
@@ -11,6 +11,13 @@ fn zero_test() {
 fn one_test() {
     assert_eq!(aliquot_sum(1), 0, "invalid result");
 }
+
+#[test]
+#[available_gas(200000)]
+fn two_test() {
+    assert_eq!(aliquot_sum(2), 1, "invalid result");
+}
+
 #[test]
 #[available_gas(200000)]
 fn one_digit_number_test() {