From 738c473b54fe81ced479e70ddbbe37515f200908 Mon Sep 17 00:00:00 2001
From: Todd Nowacki <tmn@mystenlabs.com>
Date: Tue, 11 Jun 2024 15:10:40 -0700
Subject: [PATCH] [move] Add do and range_do macros for u* (#18187)

## Description

- Added do, do_eq, range_do, and range_do_eq macros for integer types

## Test plan

- Added tests

---

## Release notes

Check each box that your changes affect. If none of the boxes relate to
your changes, release notes aren't required.

For each box you select, include information after the relevant heading
that describes the impact of your changes that a user might notice and
any actions they must take to implement updates.

- [ ] Protocol:
- [ ] Nodes (Validators and Full nodes):
- [ ] Indexer:
- [ ] JSON-RPC:
- [ ] GraphQL:
- [ ] CLI:
- [ ] Rust SDK:
---
 .../packages/move-stdlib/sources/macros.move  | 32 +++++++++++
 .../packages/move-stdlib/sources/u128.move    | 20 +++++++
 .../packages/move-stdlib/sources/u16.move     | 20 +++++++
 .../packages/move-stdlib/sources/u256.move    | 20 +++++++
 .../packages/move-stdlib/sources/u32.move     | 20 +++++++
 .../packages/move-stdlib/sources/u64.move     | 20 +++++++
 .../packages/move-stdlib/sources/u8.move      | 20 +++++++
 .../move-stdlib/tests/integer_tests.move      | 55 +++++++++++++++++++
 .../move-stdlib/tests/u128_tests.move         |  5 ++
 .../packages/move-stdlib/tests/u16_tests.move |  5 ++
 .../move-stdlib/tests/u256_tests.move         |  5 ++
 .../packages/move-stdlib/tests/u32_tests.move |  5 ++
 .../packages/move-stdlib/tests/u64_tests.move |  6 ++
 .../packages/move-stdlib/tests/u8_tests.move  |  9 +++
 14 files changed, 242 insertions(+)

diff --git a/crates/sui-framework/packages/move-stdlib/sources/macros.move b/crates/sui-framework/packages/move-stdlib/sources/macros.move
index 55116fad28b45..fb4e103cdcb17 100644
--- a/crates/sui-framework/packages/move-stdlib/sources/macros.move
+++ b/crates/sui-framework/packages/move-stdlib/sources/macros.move
@@ -68,4 +68,36 @@ module std::macros {
 
         res as $T
     }
+
+    public(package) macro fun range_do($start: _, $stop: _, $f: |_|) {
+        let mut i = $start;
+        let stop = $stop;
+        while (i < stop) {
+            $f(i);
+            i = i + 1;
+        }
+    }
+
+    public(package) macro fun range_do_eq($start: _, $stop: _, $f: |_|) {
+        let mut i = $start;
+        let stop = $stop;
+        // we check `i >= stop` inside the loop instead of `i <= stop` as `while` condition to avoid
+        // incrementing `i` past the MAX integer value.
+        // Because of this, we need to check if `i > stop` and return early--instead of letting the
+        // loop bound handle it, like in the `range_do` macro.
+        if (i > stop) return;
+        loop {
+            $f(i);
+            if (i >= stop) break;
+            i = i + 1;
+        }
+    }
+
+    public(package) macro fun do($stop: _, $f: |_|) {
+        range_do!(0, $stop, $f)
+    }
+
+    public(package) macro fun do_eq($stop: _, $f: |_|) {
+        range_do_eq!(0, $stop, $f)
+    }
 }
diff --git a/crates/sui-framework/packages/move-stdlib/sources/u128.move b/crates/sui-framework/packages/move-stdlib/sources/u128.move
index 8c4c412149ba9..947c330085ca1 100644
--- a/crates/sui-framework/packages/move-stdlib/sources/u128.move
+++ b/crates/sui-framework/packages/move-stdlib/sources/u128.move
@@ -56,4 +56,24 @@ module std::u128 {
     public fun sqrt(x: u128): u128 {
         std::macros::num_sqrt!<u128, u256>(x, 128)
     }
+
+    /// Loops applying `$f` to each number from `$start` to `$stop` (exclusive)
+    public macro fun range_do($start: u128, $stop: u128, $f: |u128|) {
+        std::macros::range_do!($start, $stop, $f)
+    }
+
+    /// Loops applying `$f` to each number from `$start` to `$stop` (inclusive)
+    public macro fun range_do_eq($start: u128, $stop: u128, $f: |u128|) {
+        std::macros::range_do_eq!($start, $stop, $f)
+    }
+
+    /// Loops applying `$f` to each number from `0` to `$stop` (exclusive)
+    public macro fun do($stop: u128, $f: |u128|) {
+        std::macros::do!($stop, $f)
+    }
+
+    /// Loops applying `$f` to each number from `0` to `$stop` (inclusive)
+    public macro fun do_eq($stop: u128, $f: |u128|) {
+        std::macros::do_eq!($stop, $f)
+    }
 }
diff --git a/crates/sui-framework/packages/move-stdlib/sources/u16.move b/crates/sui-framework/packages/move-stdlib/sources/u16.move
index 4abc5f320bc6b..9d051c11721dc 100644
--- a/crates/sui-framework/packages/move-stdlib/sources/u16.move
+++ b/crates/sui-framework/packages/move-stdlib/sources/u16.move
@@ -56,4 +56,24 @@ module std::u16 {
     public fun sqrt(x: u16): u16 {
         std::macros::num_sqrt!<u16, u32>(x, 16)
     }
+
+    /// Loops applying `$f` to each number from `$start` to `$stop` (exclusive)
+    public macro fun range_do($start: u16, $stop: u16, $f: |u16|) {
+        std::macros::range_do!($start, $stop, $f)
+    }
+
+    /// Loops applying `$f` to each number from `$start` to `$stop` (inclusive)
+    public macro fun range_do_eq($start: u16, $stop: u16, $f: |u16|) {
+        std::macros::range_do_eq!($start, $stop, $f)
+    }
+
+    /// Loops applying `$f` to each number from `0` to `$stop` (exclusive)
+    public macro fun do($stop: u16, $f: |u16|) {
+        std::macros::do!($stop, $f)
+    }
+
+    /// Loops applying `$f` to each number from `0` to `$stop` (inclusive)
+    public macro fun do_eq($stop: u16, $f: |u16|) {
+        std::macros::do_eq!($stop, $f)
+    }
 }
diff --git a/crates/sui-framework/packages/move-stdlib/sources/u256.move b/crates/sui-framework/packages/move-stdlib/sources/u256.move
index 9ec6e6b9ab0d2..1c1846db661d7 100644
--- a/crates/sui-framework/packages/move-stdlib/sources/u256.move
+++ b/crates/sui-framework/packages/move-stdlib/sources/u256.move
@@ -27,4 +27,24 @@ module std::u256 {
     public fun pow(base: u256, exponent: u8): u256 {
         std::macros::num_pow!(base, exponent)
     }
+
+    /// Loops applying `$f` to each number from `$start` to `$stop` (exclusive)
+    public macro fun range_do($start: u256, $stop: u256, $f: |u256|) {
+        std::macros::range_do!($start, $stop, $f)
+    }
+
+    /// Loops applying `$f` to each number from `$start` to `$stop` (inclusive)
+    public macro fun range_do_eq($start: u256, $stop: u256, $f: |u256|) {
+        std::macros::range_do_eq!($start, $stop, $f)
+    }
+
+    /// Loops applying `$f` to each number from `0` to `$stop` (exclusive)
+    public macro fun do($stop: u256, $f: |u256|) {
+        std::macros::do!($stop, $f)
+    }
+
+    /// Loops applying `$f` to each number from `0` to `$stop` (inclusive)
+    public macro fun do_eq($stop: u256, $f: |u256|) {
+        std::macros::do_eq!($stop, $f)
+    }
 }
diff --git a/crates/sui-framework/packages/move-stdlib/sources/u32.move b/crates/sui-framework/packages/move-stdlib/sources/u32.move
index 3b85c88aa6763..8ad44d722c178 100644
--- a/crates/sui-framework/packages/move-stdlib/sources/u32.move
+++ b/crates/sui-framework/packages/move-stdlib/sources/u32.move
@@ -56,4 +56,24 @@ module std::u32 {
     public fun sqrt(x: u32): u32 {
         std::macros::num_sqrt!<u32, u64>(x, 32)
     }
+
+    /// Loops applying `$f` to each number from `$start` to `$stop` (exclusive)
+    public macro fun range_do($start: u32, $stop: u32, $f: |u32|) {
+        std::macros::range_do!($start, $stop, $f)
+    }
+
+    /// Loops applying `$f` to each number from `$start` to `$stop` (inclusive)
+    public macro fun range_do_eq($start: u32, $stop: u32, $f: |u32|) {
+        std::macros::range_do_eq!($start, $stop, $f)
+    }
+
+    /// Loops applying `$f` to each number from `0` to `$stop` (exclusive)
+    public macro fun do($stop: u32, $f: |u32|) {
+        std::macros::do!($stop, $f)
+    }
+
+    /// Loops applying `$f` to each number from `0` to `$stop` (inclusive)
+    public macro fun do_eq($stop: u32, $f: |u32|) {
+        std::macros::do_eq!($stop, $f)
+    }
 }
diff --git a/crates/sui-framework/packages/move-stdlib/sources/u64.move b/crates/sui-framework/packages/move-stdlib/sources/u64.move
index ce3899c956695..9963dcc1b8206 100644
--- a/crates/sui-framework/packages/move-stdlib/sources/u64.move
+++ b/crates/sui-framework/packages/move-stdlib/sources/u64.move
@@ -56,4 +56,24 @@ module std::u64 {
     public fun sqrt(x: u64): u64 {
         std::macros::num_sqrt!<u64, u128>(x, 64)
     }
+
+    /// Loops applying `$f` to each number from `$start` to `$stop` (exclusive)
+    public macro fun range_do($start: u64, $stop: u64, $f: |u64|) {
+        std::macros::range_do!($start, $stop, $f)
+    }
+
+    /// Loops applying `$f` to each number from `$start` to `$stop` (inclusive)
+    public macro fun range_do_eq($start: u64, $stop: u64, $f: |u64|) {
+        std::macros::range_do_eq!($start, $stop, $f)
+    }
+
+    /// Loops applying `$f` to each number from `0` to `$stop` (exclusive)
+    public macro fun do($stop: u64, $f: |u64|) {
+        std::macros::do!($stop, $f)
+    }
+
+    /// Loops applying `$f` to each number from `0` to `$stop` (inclusive)
+    public macro fun do_eq($stop: u64, $f: |u64|) {
+        std::macros::do_eq!($stop, $f)
+    }
 }
diff --git a/crates/sui-framework/packages/move-stdlib/sources/u8.move b/crates/sui-framework/packages/move-stdlib/sources/u8.move
index 2194ac52eaecb..4eaca05a94345 100644
--- a/crates/sui-framework/packages/move-stdlib/sources/u8.move
+++ b/crates/sui-framework/packages/move-stdlib/sources/u8.move
@@ -56,4 +56,24 @@ module std::u8 {
     public fun sqrt(x: u8): u8 {
         std::macros::num_sqrt!<u8, u16>(x, 8)
     }
+
+    /// Loops applying `$f` to each number from `$start` to `$stop` (exclusive)
+    public macro fun range_do($start: u8, $stop: u8, $f: |u8|) {
+        std::macros::range_do!($start, $stop, $f)
+    }
+
+    /// Loops applying `$f` to each number from `$start` to `$stop` (inclusive)
+    public macro fun range_do_eq($start: u8, $stop: u8, $f: |u8|) {
+        std::macros::range_do_eq!($start, $stop, $f)
+    }
+
+    /// Loops applying `$f` to each number from `0` to `$stop` (exclusive)
+    public macro fun do($stop: u8, $f: |u8|) {
+        std::macros::do!($stop, $f)
+    }
+
+    /// Loops applying `$f` to each number from `0` to `$stop` (inclusive)
+    public macro fun do_eq($stop: u8, $f: |u8|) {
+        std::macros::do_eq!($stop, $f)
+    }
 }
diff --git a/crates/sui-framework/packages/move-stdlib/tests/integer_tests.move b/crates/sui-framework/packages/move-stdlib/tests/integer_tests.move
index 034a981d116a7..07d940b7a86d4 100644
--- a/crates/sui-framework/packages/move-stdlib/tests/integer_tests.move
+++ b/crates/sui-framework/packages/move-stdlib/tests/integer_tests.move
@@ -158,7 +158,62 @@ module std::integer_tests {
                 i = i + 1;
             }
         }
+    }
+
+    public(package) macro fun sum_range<$T>($n: $T): $T {
+        let n = $n;
+        (n * (n + 1)) / 2
+    }
+
+    public(package) macro fun test_dos_case<$T>($case: $T) {
+        let case = $case;
+        let mut sum: $T = 0;
+        case.do!(|i| sum = sum + i);
+        assert_eq!(sum, sum_range!(case - 1));
+
+        sum = 0;
+        case.do_eq!(|i| sum = sum + i);
+        assert_eq!(sum, sum_range!(case));
+
+        let half = case / 2;
+
+        sum = 0;
+        half.range_do!(case, |i| sum = sum + i);
+        assert_eq!(sum, sum_range!(case - 1) - sum_range!(half - 1));
+
+        sum = 0;
+        half.range_do_eq!(case, |i| sum = sum + i);
+        assert_eq!(sum, sum_range!(case) - sum_range!(half - 1));
+    }
 
+    public(package) macro fun test_dos<$T>($max: $T, $cases: vector<$T>) {
+        let max = $max;
+        let cases = $cases;
+        // test bounds/invalid ranges
+        (0: $T).do!(|_| assert!(false));
+        cases!(max, cases, |case_pred, case, case_succ| {
+            if (case == 0) return;
+            case.range_do!(0, |_| assert!(false));
+            case.range_do_eq!(0, |_| assert!(false));
+
+            if (case == max) return;
+            case.range_do!(case_pred, |_| assert!(false));
+            case_succ.range_do!(case, |_| assert!(false));
+            case.range_do_eq!(case_pred, |_| assert!(false));
+            case_succ.range_do_eq!(case, |_| assert!(false));
+        });
+
+        // test upper bound being max
+        let max_pred = max - 1;
+        max_pred.range_do_eq!(max, |_| ());
+
+        // test iteration numbers
+        let cases: vector<$T> = vector[3, 5, 8, 11, 14];
+        cases!(max, cases, |case_pred, case, case_succ| {
+            test_dos_case!(case_pred);
+            test_dos_case!(case);
+            test_dos_case!(case_succ);
+        });
     }
 
 
diff --git a/crates/sui-framework/packages/move-stdlib/tests/u128_tests.move b/crates/sui-framework/packages/move-stdlib/tests/u128_tests.move
index fc433df52e81d..315190be0589b 100644
--- a/crates/sui-framework/packages/move-stdlib/tests/u128_tests.move
+++ b/crates/sui-framework/packages/move-stdlib/tests/u128_tests.move
@@ -70,4 +70,9 @@ module std::u128_tests {
             vector[0, 2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35, 38, 41, 44, 47, 50, 53, 56, 59];
         integer_tests::test_sqrt!(MAX, CASES, reflexive_cases)
     }
+
+    #[test]
+    fun test_dos() {
+        integer_tests::test_dos!(MAX, CASES);
+    }
 }
diff --git a/crates/sui-framework/packages/move-stdlib/tests/u16_tests.move b/crates/sui-framework/packages/move-stdlib/tests/u16_tests.move
index dac3c4340cfec..f872acf70f988 100644
--- a/crates/sui-framework/packages/move-stdlib/tests/u16_tests.move
+++ b/crates/sui-framework/packages/move-stdlib/tests/u16_tests.move
@@ -70,4 +70,9 @@ module std::u16_tests {
             vector[0, 2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35, 38, 41, 44, 47, 50, 53, 56, 59];
         integer_tests::test_sqrt!(MAX, CASES, reflexive_cases)
     }
+
+    #[test]
+    fun test_dos() {
+        integer_tests::test_dos!(MAX, CASES);
+    }
 }
diff --git a/crates/sui-framework/packages/move-stdlib/tests/u256_tests.move b/crates/sui-framework/packages/move-stdlib/tests/u256_tests.move
index 692ab83f601fb..eff171bf43fa0 100644
--- a/crates/sui-framework/packages/move-stdlib/tests/u256_tests.move
+++ b/crates/sui-framework/packages/move-stdlib/tests/u256_tests.move
@@ -64,4 +64,9 @@ module std::u256_tests {
     fun test_pow_overflow() {
         255u256.pow(255);
     }
+
+    #[test]
+    fun test_dos() {
+        integer_tests::test_dos!(MAX, CASES);
+    }
 }
diff --git a/crates/sui-framework/packages/move-stdlib/tests/u32_tests.move b/crates/sui-framework/packages/move-stdlib/tests/u32_tests.move
index 5ec4db28bc924..c2890c1d5cff9 100644
--- a/crates/sui-framework/packages/move-stdlib/tests/u32_tests.move
+++ b/crates/sui-framework/packages/move-stdlib/tests/u32_tests.move
@@ -70,4 +70,9 @@ module std::u32_tests {
             vector[0, 2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35, 38, 41, 44, 47, 50, 53, 56, 59];
         integer_tests::test_sqrt!(MAX, CASES, reflexive_cases)
     }
+
+    #[test]
+    fun test_dos() {
+        integer_tests::test_dos!(MAX, CASES);
+    }
 }
diff --git a/crates/sui-framework/packages/move-stdlib/tests/u64_tests.move b/crates/sui-framework/packages/move-stdlib/tests/u64_tests.move
index e598991eea77f..aee0790856f0f 100644
--- a/crates/sui-framework/packages/move-stdlib/tests/u64_tests.move
+++ b/crates/sui-framework/packages/move-stdlib/tests/u64_tests.move
@@ -70,4 +70,10 @@ module std::u64_tests {
             vector[0, 2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35, 38, 41, 44, 47, 50, 53, 56, 59];
         integer_tests::test_sqrt!(MAX, CASES, reflexive_cases)
     }
+
+    #[test]
+    fun test_dos() {
+        integer_tests::test_dos!(MAX, CASES);
+    }
+
 }
diff --git a/crates/sui-framework/packages/move-stdlib/tests/u8_tests.move b/crates/sui-framework/packages/move-stdlib/tests/u8_tests.move
index 8f8c9b183c279..7aa1d6f263a00 100644
--- a/crates/sui-framework/packages/move-stdlib/tests/u8_tests.move
+++ b/crates/sui-framework/packages/move-stdlib/tests/u8_tests.move
@@ -4,6 +4,7 @@
 #[test_only]
 module std::u8_tests {
     use std::integer_tests;
+    use std::unit_test::assert_eq;
 
     const BIT_SIZE: u8 = 8;
     const MAX: u8 = 0xFF;
@@ -65,4 +66,12 @@ module std::u8_tests {
     fun test_sqrt() {
         integer_tests::test_sqrt!(MAX, CASES, vector[0, 2, 5, 8, 11, 14]);
     }
+
+    #[test]
+    fun test_dos() {
+        let mut sum = 0u16;
+        255u8.do_eq!(|i| sum = sum + (i as u16));
+        assert_eq!(sum, 32640);
+        integer_tests::test_dos!(MAX, CASES);
+    }
 }