feat: update aliquot_sum (#309)

optimize code and add test two
keep-starknet-strange · Jun 24, 2024 · 0aa7cca · 0aa7cca
1 parent ef3292b
commit 0aa7cca
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 7 deletions.
diff --git 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
@@ -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() {