Optimizing vector::replace and compiler-v2 vector depenency fix (#15524)

* Optimizing vector::replace with mem::replace * [compiler-v2] Also include all dependencies of vector for compiler v2 from #15528 --------- Co-authored-by: Igor <[email protected]>
aptos-labs · Dec 19, 2024 · 1381c93 · 1381c93
1 parent 127df5c
commit 1381c93
Show file tree

Hide file tree

Showing 5 changed files with 42 additions and 26 deletions.
diff --git a/aptos-move/framework/aptos-stdlib/doc/big_vector.md b/aptos-move/framework/aptos-stdlib/doc/big_vector.md
@@ -502,7 +502,7 @@ Aborts if <code>i</code> is out of bounds.
     <b>if</b> (self.end_index == i) {
         <b>return</b> last_val
     };
-    // because the lack of mem::swap, here we swap remove the requested value from the bucket
+    // because the lack of <a href="../../move-stdlib/doc/mem.md#0x1_mem_swap">mem::swap</a>, here we swap remove the requested value from the bucket
     // and append the last_val <b>to</b> the bucket then swap the last bucket val back
     <b>let</b> bucket = <a href="table_with_length.md#0x1_table_with_length_borrow_mut">table_with_length::borrow_mut</a>(&<b>mut</b> self.buckets, i / self.bucket_size);
     <b>let</b> bucket_len = <a href="../../move-stdlib/doc/vector.md#0x1_vector_length">vector::length</a>(bucket);

diff --git a/aptos-move/framework/move-stdlib/doc/vector.md b/aptos-move/framework/move-stdlib/doc/vector.md
@@ -88,7 +88,8 @@ the return on investment didn't seem worth it for these simple functions.
     -  [Function `rotate_slice`](#@Specification_1_rotate_slice)
 
 
-<pre><code></code></pre>
+<pre><code><b>use</b> <a href="mem.md#0x1_mem">0x1::mem</a>;
+</code></pre>
 
 
 
@@ -930,14 +931,13 @@ Aborts if <code>i</code> is out of bounds.
 <pre><code><b>public</b> <b>fun</b> <a href="vector.md#0x1_vector_replace">replace</a>&lt;Element&gt;(self: &<b>mut</b> <a href="vector.md#0x1_vector">vector</a>&lt;Element&gt;, i: u64, val: Element): Element {
     <b>let</b> last_idx = <a href="vector.md#0x1_vector_length">length</a>(self);
     <b>assert</b>!(i &lt; last_idx, <a href="vector.md#0x1_vector_EINDEX_OUT_OF_BOUNDS">EINDEX_OUT_OF_BOUNDS</a>);
-    // TODO: Enable after tests are fixed.
-    // <b>if</b> (<a href="vector.md#0x1_vector_USE_MOVE_RANGE">USE_MOVE_RANGE</a>) {
-    //     <a href="mem.md#0x1_mem_replace">mem::replace</a>(<a href="vector.md#0x1_vector_borrow_mut">borrow_mut</a>(self, i), val)
-    // } <b>else</b> {
-    <a href="vector.md#0x1_vector_push_back">push_back</a>(self, val);
-    <a href="vector.md#0x1_vector_swap">swap</a>(self, i, last_idx);
-    <a href="vector.md#0x1_vector_pop_back">pop_back</a>(self)
-    // }
+    <b>if</b> (<a href="vector.md#0x1_vector_USE_MOVE_RANGE">USE_MOVE_RANGE</a>) {
+        <a href="mem.md#0x1_mem_replace">mem::replace</a>(<a href="vector.md#0x1_vector_borrow_mut">borrow_mut</a>(self, i), val)
+    } <b>else</b> {
+        <a href="vector.md#0x1_vector_push_back">push_back</a>(self, val);
+        <a href="vector.md#0x1_vector_swap">swap</a>(self, i, last_idx);
+        <a href="vector.md#0x1_vector_pop_back">pop_back</a>(self)
+    }
 }
 </code></pre>
 

diff --git a/aptos-move/framework/move-stdlib/sources/mem.move b/aptos-move/framework/move-stdlib/sources/mem.move
@@ -2,7 +2,7 @@
 module std::mem {
     // TODO - functions here are `public(friend)` here for one release,
     // and to be changed to `public` one release later.
-    // friend std::vector;
+    friend std::vector;
     #[test_only]
     friend std::mem_tests;
 

diff --git a/aptos-move/framework/move-stdlib/sources/vector.move b/aptos-move/framework/move-stdlib/sources/vector.move
@@ -9,7 +9,7 @@
 /// Move functions here because many have loops, requiring loop invariants to prove, and
 /// the return on investment didn't seem worth it for these simple functions.
 module std::vector {
-    // use std::mem;
+    use std::mem;
 
     /// The index into the vector is out of bounds
     const EINDEX_OUT_OF_BOUNDS: u64 = 0x20000;
@@ -357,14 +357,13 @@ module std::vector {
     public fun replace<Element>(self: &mut vector<Element>, i: u64, val: Element): Element {
         let last_idx = length(self);
         assert!(i < last_idx, EINDEX_OUT_OF_BOUNDS);
-        // TODO: Enable after tests are fixed.
-        // if (USE_MOVE_RANGE) {
-        //     mem::replace(borrow_mut(self, i), val)
-        // } else {
-        push_back(self, val);
-        swap(self, i, last_idx);
-        pop_back(self)
-        // }
+        if (USE_MOVE_RANGE) {
+            mem::replace(borrow_mut(self, i), val)
+        } else {
+            push_back(self, val);
+            swap(self, i, last_idx);
+            pop_back(self)
+        }
     }
 
     /// Apply the function to each element in the vector, consuming it.

diff --git a/third_party/move/move-model/src/lib.rs b/third_party/move/move-model/src/lib.rs
@@ -305,13 +305,25 @@ pub fn run_model_builder_with_options_and_compilation_flags<
         },
     };
 
-    // Extract the module/script closure
+    // Extract the module/script dependency closure
     let mut visited_modules = BTreeSet::new();
+    // Extract the module dependency closure for the vector module
+    let mut vector_and_its_dependencies = BTreeSet::new();
+    let mut seen_vector = false;
     for (_, mident, mdef) in &expansion_ast.modules {
         let src_file_hash = mdef.loc.file_hash();
         if !dep_files.contains(&src_file_hash) {
             collect_related_modules_recursive(mident, &expansion_ast.modules, &mut visited_modules);
         }
+        if !seen_vector && is_vector(*mident) {
+            seen_vector = true;
+            // Collect the vector module and its dependencies.
+            collect_related_modules_recursive(
+                mident,
+                &expansion_ast.modules,
+                &mut vector_and_its_dependencies,
+            );
+        }
     }
     for sdef in expansion_ast.scripts.values() {
         let src_file_hash = sdef.loc.file_hash();
@@ -330,14 +342,13 @@ pub fn run_model_builder_with_options_and_compilation_flags<
     let mut expansion_ast = {
         let E::Program { modules, scripts } = expansion_ast;
         let modules = modules.filter_map(|mident, mut mdef| {
-            // We need to always include the `vector` module (only for compiler v2),
+            // For compiler v2, we need to always include the `vector` module and any of its dependencies,
             // to handle cases of implicit usage.
             // E.g., index operation on a vector results in a call to `vector::borrow`.
             // TODO(#15483): consider refactoring code to avoid this special case.
-            let is_vector = mident.value.address.into_addr_bytes().into_inner()
-                == AccountAddress::ONE
-                && mident.value.module.0.value.as_str() == "vector";
-            (is_vector && compile_via_model || visited_modules.contains(&mident.value)).then(|| {
+            ((compile_via_model && vector_and_its_dependencies.contains(&mident.value))
+                || visited_modules.contains(&mident.value))
+            .then(|| {
                 mdef.is_source_module = true;
                 mdef
             })
@@ -416,6 +427,12 @@ pub fn run_model_builder_with_options_and_compilation_flags<
     }
 }
 
+/// Is `module_ident` the `vector` module?
+fn is_vector(module_ident: ModuleIdent_) -> bool {
+    module_ident.address.into_addr_bytes().into_inner() == AccountAddress::ONE
+        && module_ident.module.0.value.as_str() == "vector"
+}
+
 fn run_move_checker(env: &mut GlobalEnv, program: E::Program) {
     let mut builder = ModelBuilder::new(env);
     for (module_count, (module_id, module_def)) in program