Add range function in vector module (#11173)

GitOrigin-RevId: a59ae0a3e198be5cb29cfd77b8abda842989e728
aptos-labs · Oct 23, 2024 · bd6a824 · bd6a824
1 parent ad32817
commit bd6a824
Show file tree

Hide file tree

Showing 3 changed files with 207 additions and 0 deletions.
diff --git a/move-stdlib/doc/vector.md b/move-stdlib/doc/vector.md
@@ -63,6 +63,9 @@ the return on investment didn't seem worth it for these simple functions.
 -  [Function `any`](#0x1_vector_any)
 -  [Function `all`](#0x1_vector_all)
 -  [Function `destroy`](#0x1_vector_destroy)
+-  [Function `range`](#0x1_vector_range)
+-  [Function `range_with_step`](#0x1_vector_range_with_step)
+-  [Function `slice`](#0x1_vector_slice)
 -  [Specification](#@Specification_1)
     -  [Helper Functions](#@Helper_Functions_2)
     -  [Function `singleton`](#@Specification_1_singleton)
@@ -112,6 +115,26 @@ The index into the vector is out of bounds
 
 
 
+<a id="0x1_vector_EINVALID_SLICE_RANGE"></a>
+
+The range in <code>slice</code> is invalid.
+
+
+<pre><code><b>const</b> <a href="vector.md#0x1_vector_EINVALID_SLICE_RANGE">EINVALID_SLICE_RANGE</a>: u64 = 131076;
+</code></pre>
+
+
+
+<a id="0x1_vector_EINVALID_STEP"></a>
+
+The step provided in <code>range</code> is invalid, must be greater than zero.
+
+
+<pre><code><b>const</b> <a href="vector.md#0x1_vector_EINVALID_STEP">EINVALID_STEP</a>: u64 = 131075;
+</code></pre>
+
+
+
 <a id="0x1_vector_EVECTORS_LENGTH_MISMATCH"></a>
 
 The length of the vectors are not equal.
@@ -1580,6 +1603,96 @@ when used in the context of destroying a vector.
 
 
 
+</details>
+
+<a id="0x1_vector_range"></a>
+
+## Function `range`
+
+
+
+<pre><code><b>public</b> <b>fun</b> <a href="vector.md#0x1_vector_range">range</a>(start: u64, end: u64): <a href="vector.md#0x1_vector">vector</a>&lt;u64&gt;
+</code></pre>
+
+
+
+<details>
+<summary>Implementation</summary>
+
+
+<pre><code><b>public</b> <b>fun</b> <a href="vector.md#0x1_vector_range">range</a>(start: u64, end: u64): <a href="vector.md#0x1_vector">vector</a>&lt;u64&gt; {
+    <a href="vector.md#0x1_vector_range_with_step">range_with_step</a>(start, end, 1)
+}
+</code></pre>
+
+
+
+</details>
+
+<a id="0x1_vector_range_with_step"></a>
+
+## Function `range_with_step`
+
+
+
+<pre><code><b>public</b> <b>fun</b> <a href="vector.md#0x1_vector_range_with_step">range_with_step</a>(start: u64, end: u64, step: u64): <a href="vector.md#0x1_vector">vector</a>&lt;u64&gt;
+</code></pre>
+
+
+
+<details>
+<summary>Implementation</summary>
+
+
+<pre><code><b>public</b> <b>fun</b> <a href="vector.md#0x1_vector_range_with_step">range_with_step</a>(start: u64, end: u64, step: u64): <a href="vector.md#0x1_vector">vector</a>&lt;u64&gt; {
+    <b>assert</b>!(step &gt; 0, <a href="vector.md#0x1_vector_EINVALID_STEP">EINVALID_STEP</a>);
+
+    <b>let</b> vec = <a href="vector.md#0x1_vector">vector</a>[];
+    <b>while</b> (start &lt; end) {
+        <a href="vector.md#0x1_vector_push_back">push_back</a>(&<b>mut</b> vec, start);
+        start = start + step;
+    };
+    vec
+}
+</code></pre>
+
+
+
+</details>
+
+<a id="0x1_vector_slice"></a>
+
+## Function `slice`
+
+
+
+<pre><code><b>public</b> <b>fun</b> <a href="vector.md#0x1_vector_slice">slice</a>&lt;Element: <b>copy</b>&gt;(v: &<a href="vector.md#0x1_vector">vector</a>&lt;Element&gt;, start: u64, end: u64): <a href="vector.md#0x1_vector">vector</a>&lt;Element&gt;
+</code></pre>
+
+
+
+<details>
+<summary>Implementation</summary>
+
+
+<pre><code><b>public</b> <b>fun</b> <a href="vector.md#0x1_vector_slice">slice</a>&lt;Element: <b>copy</b>&gt;(
+    v: &<a href="vector.md#0x1_vector">vector</a>&lt;Element&gt;,
+    start: u64,
+    end: u64
+): <a href="vector.md#0x1_vector">vector</a>&lt;Element&gt; {
+    <b>assert</b>!(start &lt;= end && end &lt;= <a href="vector.md#0x1_vector_length">length</a>(v), <a href="vector.md#0x1_vector_EINVALID_SLICE_RANGE">EINVALID_SLICE_RANGE</a>);
+
+    <b>let</b> vec = <a href="vector.md#0x1_vector">vector</a>[];
+    <b>while</b> (start &lt; end) {
+        <a href="vector.md#0x1_vector_push_back">push_back</a>(&<b>mut</b> vec, *<a href="vector.md#0x1_vector_borrow">borrow</a>(v, start));
+        start = start + 1;
+    };
+    vec
+}
+</code></pre>
+
+
+
 </details>
 
 <a id="@Specification_1"></a>

diff --git a/move-stdlib/sources/vector.move b/move-stdlib/sources/vector.move
@@ -18,6 +18,12 @@ module std::vector {
     /// The length of the vectors are not equal.
     const EVECTORS_LENGTH_MISMATCH: u64 = 0x20002;
 
+    /// The step provided in `range` is invalid, must be greater than zero.
+    const EINVALID_STEP: u64 = 0x20003;
+
+    /// The range in `slice` is invalid.
+    const EINVALID_SLICE_RANGE: u64 = 0x20004;
+
     #[bytecode_instruction]
     /// Create an empty vector.
     native public fun empty<Element>(): vector<Element>;
@@ -589,6 +595,36 @@ module std::vector {
         for_each_reverse(v, |e| d(e))
     }
 
+    public fun range(start: u64, end: u64): vector<u64> {
+        range_with_step(start, end, 1)
+    }
+
+    public fun range_with_step(start: u64, end: u64, step: u64): vector<u64> {
+        assert!(step > 0, EINVALID_STEP);
+
+        let vec = vector[];
+        while (start < end) {
+            push_back(&mut vec, start);
+            start = start + step;
+        };
+        vec
+    }
+
+    public fun slice<Element: copy>(
+        v: &vector<Element>,
+        start: u64,
+        end: u64
+    ): vector<Element> {
+        assert!(start <= end && end <= length(v), EINVALID_SLICE_RANGE);
+
+        let vec = vector[];
+        while (start < end) {
+            push_back(&mut vec, *borrow(v, start));
+            start = start + 1;
+        };
+        vec
+    }
+
     // =================================================================
     // Module Specification
 

diff --git a/move-stdlib/tests/vector_tests.move b/move-stdlib/tests/vector_tests.move
@@ -888,6 +888,64 @@ module std::vector_tests {
         vector::insert(&mut v,6, 6);
     }
 
+    #[test]
+    fun test_range() {
+        let result = vector::range(5, 10);
+        assert!(result == vector[5, 6, 7, 8, 9], 1);
+    }
+
+    #[test]
+    fun test_range_with_step() {
+        let result = vector::range_with_step(0, 10, 2);
+        assert!(result == vector[0, 2, 4, 6, 8], 1);
+
+        let empty_result = vector::range_with_step(10, 10, 2);
+        assert!(empty_result == vector[], 1);
+
+        // Test with `start` greater than `end`
+        let reverse_result = vector::range_with_step(10, 0, 2);
+        assert!(reverse_result == vector[], 1);
+    }
+
+    #[test]
+    #[expected_failure(abort_code = V::EINVALID_STEP)]
+    fun test_range_with_invalid_step() {
+        vector::range_with_step(0, 10, 0);
+    }
+
+    #[test]
+    fun test_slice() {
+        let v = &vector[0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
+
+        let slice_beginning = vector::slice(v, 0, 3);
+        assert!(slice_beginning == vector[0, 1, 2], 1);
+
+        let slice_end = vector::slice(v, 7, 10);
+        assert!(slice_end == vector[7, 8, 9], 1);
+
+        let empty_slice = vector::slice(v, 5, 5);
+        assert!(empty_slice == vector[], 1);
+        let empty_slice = vector::slice(v, 0, 0);
+        assert!(empty_slice == vector[], 1);
+
+        let full_slice = &vector::slice(v, 0, 10);
+        assert!(full_slice == v, 1);
+    }
+
+    #[test]
+    #[expected_failure(abort_code = V::EINVALID_SLICE_RANGE)]
+    fun test_slice_invalid_range() {
+        let v = &vector[0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
+        vector::slice(v, 7, 6); // start is greater than end
+    }
+
+    #[test]
+    #[expected_failure(abort_code = V::EINVALID_SLICE_RANGE)]
+    fun test_slice_out_of_bounds() {
+        let v = &vector[0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
+        vector::slice(v, 0, 11); // end is out of bounds
+    }
+
     #[test_only]
     struct MoveOnly {}