adapter(visitor): Use UIDTraversal across all execution layers (#19576)

## Description Replace the legacy implementations of `get_all_uids` with the one that has been used in `latest` for some time now. This version avoids inflating the whole struct just to read its UIDs. The main motivation in making the switch is to reduce the number of call-sites that use `simple_deserialize`, so that we can eventually get rid of it entirely, and replace it with a visitor based ipmplementation that includes the appropriate caveats around inflating a whole struct in memory. ## Test plan ``` sui$ cargo simtest ``` ## Stack - #19517 - #19518 - #19519 - #19520 - #19521 - #19548 - #19554 - #19565 --- ## 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: - [ ] REST API: --------- Co-authored-by: Timothy Zakian <tim@mystenlabs.com>
MystenLabs · Oct 1, 2024 · b7333cd · b7333cd
1 parent fd81820
commit b7333cd
Show file tree

Hide file tree

Showing 4 changed files with 114 additions and 111 deletions.
diff --git a/sui-execution/latest/sui-move-natives/src/object_runtime/mod.rs b/sui-execution/latest/sui-move-natives/src/object_runtime/mod.rs
@@ -690,26 +690,26 @@ pub fn get_all_uids(
     bcs_bytes: &[u8],
 ) -> Result<BTreeSet<ObjectID>, /* invariant violation */ String> {
     let mut ids = BTreeSet::new();
-    struct UIDTraversalV2<'i>(&'i mut BTreeSet<ObjectID>);
-    struct UIDCollectorV2<'i>(&'i mut BTreeSet<ObjectID>);
+    struct UIDTraversal<'i>(&'i mut BTreeSet<ObjectID>);
+    struct UIDCollector<'i>(&'i mut BTreeSet<ObjectID>);
 
-    impl<'i, 'b, 'l> AV::Traversal<'b, 'l> for UIDTraversalV2<'i> {
+    impl<'i, 'b, 'l> AV::Traversal<'b, 'l> for UIDTraversal<'i> {
         type Error = AV::Error;
 
         fn traverse_struct(
             &mut self,
             driver: &mut AV::StructDriver<'_, 'b, 'l>,
         ) -> Result<(), Self::Error> {
             if driver.struct_layout().type_ == UID::type_() {
-                while driver.next_field(&mut UIDCollectorV2(self.0))?.is_some() {}
+                while driver.next_field(&mut UIDCollector(self.0))?.is_some() {}
             } else {
                 while driver.next_field(self)?.is_some() {}
             }
             Ok(())
         }
     }
 
-    impl<'i, 'b, 'l> AV::Traversal<'b, 'l> for UIDCollectorV2<'i> {
+    impl<'i, 'b, 'l> AV::Traversal<'b, 'l> for UIDCollector<'i> {
         type Error = AV::Error;
         fn traverse_address(
             &mut self,
@@ -724,7 +724,7 @@ pub fn get_all_uids(
     MoveValue::visit_deserialize(
         bcs_bytes,
         fully_annotated_layout,
-        &mut UIDTraversalV2(&mut ids),
+        &mut UIDTraversal(&mut ids),
     )
     .map_err(|e| format!("Failed to deserialize. {e:?}"))?;
     Ok(ids)

diff --git a/sui-execution/v0/sui-move-natives/src/object_runtime/mod.rs b/sui-execution/v0/sui-move-natives/src/object_runtime/mod.rs
@@ -5,7 +5,7 @@ use better_any::{Tid, TidAble};
 use linked_hash_map::LinkedHashMap;
 use move_binary_format::errors::{PartialVMError, PartialVMResult};
 use move_core_types::{
-    account_address::AccountAddress, annotated_value as A, effects::Op,
+    account_address::AccountAddress, annotated_value as A, annotated_visitor as AV, effects::Op,
     language_storage::StructTag, runtime_value as R, vm_status::StatusCode,
 };
 use move_vm_types::{
@@ -608,7 +608,6 @@ fn update_owner_map(
     Ok(())
 }
 
-// TODO use a custom DeserializerSeed and improve this performance
 /// WARNING! This function assumes that the bcs bytes have already been validated,
 /// and it will give an invariant violation otherwise.
 /// In short, we are relying on the invariant that the bytes are valid for objects
@@ -619,40 +618,42 @@ pub fn get_all_uids(
     bcs_bytes: &[u8],
 ) -> Result<BTreeSet<ObjectID>, /* invariant violation */ String> {
     let mut ids = BTreeSet::new();
-    // TODO (annotated-visitor): Replace with a custom visitor
-    let v = A::MoveValue::simple_deserialize(bcs_bytes, fully_annotated_layout)
-        .map_err(|e| format!("Failed to deserialize. {e:?}"))?;
-    get_all_uids_in_value(&mut ids, &v)?;
-    Ok(ids)
-}
-
-fn get_all_uids_in_value(
-    acc: &mut BTreeSet<ObjectID>,
-    v: &A::MoveValue,
-) -> Result<(), /* invariant violation */ String> {
-    let mut stack = vec![v];
-    while let Some(cur) = stack.pop() {
-        let s = match cur {
-            A::MoveValue::Struct(s) => s,
-            A::MoveValue::Vector(vec) => {
-                stack.extend(vec);
-                continue;
+    struct UIDTraversal<'i>(&'i mut BTreeSet<ObjectID>);
+    struct UIDCollector<'i>(&'i mut BTreeSet<ObjectID>);
+
+    impl<'i, 'b, 'l> AV::Traversal<'b, 'l> for UIDTraversal<'i> {
+        type Error = AV::Error;
+
+        fn traverse_struct(
+            &mut self,
+            driver: &mut AV::StructDriver<'_, 'b, 'l>,
+        ) -> Result<(), Self::Error> {
+            if driver.struct_layout().type_ == UID::type_() {
+                while driver.next_field(&mut UIDCollector(self.0))?.is_some() {}
+            } else {
+                while driver.next_field(self)?.is_some() {}
             }
-            _ => continue,
-        };
-        let A::MoveStruct { type_, fields } = s;
-        if type_ == &UID::type_() {
-            let inner = match &fields[0].1 {
-                A::MoveValue::Struct(A::MoveStruct { fields, .. }) => fields,
-                v => return Err(format!("Unexpected UID layout. {v:?}")),
-            };
-            match &inner[0].1 {
-                A::MoveValue::Address(id) => acc.insert((*id).into()),
-                v => return Err(format!("Unexpected ID layout. {v:?}")),
-            };
-        } else {
-            stack.extend(fields.iter().map(|(_, v)| v));
+            Ok(())
         }
     }
-    Ok(())
+
+    impl<'i, 'b, 'l> AV::Traversal<'b, 'l> for UIDCollector<'i> {
+        type Error = AV::Error;
+        fn traverse_address(
+            &mut self,
+            _driver: &AV::ValueDriver<'_, 'b, 'l>,
+            value: AccountAddress,
+        ) -> Result<(), Self::Error> {
+            self.0.insert(value.into());
+            Ok(())
+        }
+    }
+
+    A::MoveValue::visit_deserialize(
+        bcs_bytes,
+        fully_annotated_layout,
+        &mut UIDTraversal(&mut ids),
+    )
+    .map_err(|e| format!("Failed to deserialize. {e:?}"))?;
+    Ok(ids)
 }
diff --git a/sui-execution/v1/sui-move-natives/src/object_runtime/mod.rs b/sui-execution/v1/sui-move-natives/src/object_runtime/mod.rs
@@ -6,7 +6,8 @@ use linked_hash_map::LinkedHashMap;
 use move_binary_format::errors::{PartialVMError, PartialVMResult};
 use move_core_types::{
     account_address::AccountAddress,
-    annotated_value::{MoveStruct, MoveTypeLayout, MoveValue},
+    annotated_value::{MoveTypeLayout, MoveValue},
+    annotated_visitor as AV,
     effects::Op,
     language_storage::StructTag,
     runtime_value as R,
@@ -660,7 +661,6 @@ fn check_circular_ownership(
     Ok(())
 }
 
-// TODO use a custom DeserializerSeed and improve this performance
 /// WARNING! This function assumes that the bcs bytes have already been validated,
 /// and it will give an invariant violation otherwise.
 /// In short, we are relying on the invariant that the bytes are valid for objects
@@ -671,41 +671,42 @@ pub fn get_all_uids(
     bcs_bytes: &[u8],
 ) -> Result<BTreeSet<ObjectID>, /* invariant violation */ String> {
     let mut ids = BTreeSet::new();
-    // TODO (annotated-visitor): Replace with a custom visitor
-    let v = MoveValue::simple_deserialize(bcs_bytes, fully_annotated_layout)
-        .map_err(|e| format!("Failed to deserialize. {e:?}"))?;
-    get_all_uids_in_value(&mut ids, &v)?;
-    Ok(ids)
-}
-
-fn get_all_uids_in_value(
-    acc: &mut BTreeSet<ObjectID>,
-    v: &MoveValue,
-) -> Result<(), /* invariant violation */ String> {
-    let mut stack = vec![v];
-    while let Some(cur) = stack.pop() {
-        let s = match cur {
-            MoveValue::Struct(s) => s,
-            MoveValue::Vector(vec) => {
-                stack.extend(vec);
-                continue;
+    struct UIDTraversal<'i>(&'i mut BTreeSet<ObjectID>);
+    struct UIDCollector<'i>(&'i mut BTreeSet<ObjectID>);
+
+    impl<'i, 'b, 'l> AV::Traversal<'b, 'l> for UIDTraversal<'i> {
+        type Error = AV::Error;
+
+        fn traverse_struct(
+            &mut self,
+            driver: &mut AV::StructDriver<'_, 'b, 'l>,
+        ) -> Result<(), Self::Error> {
+            if driver.struct_layout().type_ == UID::type_() {
+                while driver.next_field(&mut UIDCollector(self.0))?.is_some() {}
+            } else {
+                while driver.next_field(self)?.is_some() {}
             }
-            _ => continue,
-        };
+            Ok(())
+        }
+    }
 
-        let MoveStruct { type_, fields } = s;
-        if type_ == &UID::type_() {
-            let inner = match &fields[0].1 {
-                MoveValue::Struct(MoveStruct { fields, .. }) => fields,
-                v => return Err(format!("Unexpected UID layout. {v:?}")),
-            };
-            match &inner[0].1 {
-                MoveValue::Address(id) => acc.insert((*id).into()),
-                v => return Err(format!("Unexpected ID layout. {v:?}")),
-            };
-        } else {
-            stack.extend(fields.iter().map(|(_, v)| v));
+    impl<'i, 'b, 'l> AV::Traversal<'b, 'l> for UIDCollector<'i> {
+        type Error = AV::Error;
+        fn traverse_address(
+            &mut self,
+            _driver: &AV::ValueDriver<'_, 'b, 'l>,
+            value: AccountAddress,
+        ) -> Result<(), Self::Error> {
+            self.0.insert(value.into());
+            Ok(())
         }
     }
-    Ok(())
+
+    MoveValue::visit_deserialize(
+        bcs_bytes,
+        fully_annotated_layout,
+        &mut UIDTraversal(&mut ids),
+    )
+    .map_err(|e| format!("Failed to deserialize. {e:?}"))?;
+    Ok(ids)
 }
diff --git a/sui-execution/v2/sui-move-natives/src/object_runtime/mod.rs b/sui-execution/v2/sui-move-natives/src/object_runtime/mod.rs
@@ -11,7 +11,8 @@ use indexmap::set::IndexSet;
 use move_binary_format::errors::{PartialVMError, PartialVMResult};
 use move_core_types::{
     account_address::AccountAddress,
-    annotated_value::{MoveStruct, MoveTypeLayout, MoveValue},
+    annotated_value::{MoveTypeLayout, MoveValue},
+    annotated_visitor as AV,
     effects::Op,
     language_storage::StructTag,
     runtime_value as R,
@@ -628,7 +629,6 @@ fn check_circular_ownership(
     Ok(())
 }
 
-// TODO use a custom DeserializerSeed and improve this performance
 /// WARNING! This function assumes that the bcs bytes have already been validated,
 /// and it will give an invariant violation otherwise.
 /// In short, we are relying on the invariant that the bytes are valid for objects
@@ -639,41 +639,42 @@ pub fn get_all_uids(
     bcs_bytes: &[u8],
 ) -> Result<BTreeSet<ObjectID>, /* invariant violation */ String> {
     let mut ids = BTreeSet::new();
-    // TODO (annotated-visitor): Replace with a custom visitor
-    let v = MoveValue::simple_deserialize(bcs_bytes, fully_annotated_layout)
-        .map_err(|e| format!("Failed to deserialize. {e:?}"))?;
-    get_all_uids_in_value(&mut ids, &v)?;
-    Ok(ids)
-}
-
-fn get_all_uids_in_value(
-    acc: &mut BTreeSet<ObjectID>,
-    v: &MoveValue,
-) -> Result<(), /* invariant violation */ String> {
-    let mut stack = vec![v];
-    while let Some(cur) = stack.pop() {
-        let s = match cur {
-            MoveValue::Struct(s) => s,
-            MoveValue::Vector(vec) => {
-                stack.extend(vec);
-                continue;
+    struct UIDTraversal<'i>(&'i mut BTreeSet<ObjectID>);
+    struct UIDCollector<'i>(&'i mut BTreeSet<ObjectID>);
+
+    impl<'i, 'b, 'l> AV::Traversal<'b, 'l> for UIDTraversal<'i> {
+        type Error = AV::Error;
+
+        fn traverse_struct(
+            &mut self,
+            driver: &mut AV::StructDriver<'_, 'b, 'l>,
+        ) -> Result<(), Self::Error> {
+            if driver.struct_layout().type_ == UID::type_() {
+                while driver.next_field(&mut UIDCollector(self.0))?.is_some() {}
+            } else {
+                while driver.next_field(self)?.is_some() {}
             }
-            _ => continue,
-        };
+            Ok(())
+        }
+    }
 
-        let MoveStruct { type_, fields } = s;
-        if type_ == &UID::type_() {
-            let inner = match &fields[0].1 {
-                MoveValue::Struct(MoveStruct { fields, .. }) => fields,
-                v => return Err(format!("Unexpected UID layout. {v:?}")),
-            };
-            match &inner[0].1 {
-                MoveValue::Address(id) => acc.insert((*id).into()),
-                v => return Err(format!("Unexpected ID layout. {v:?}")),
-            };
-        } else {
-            stack.extend(fields.iter().map(|(_, v)| v));
+    impl<'i, 'b, 'l> AV::Traversal<'b, 'l> for UIDCollector<'i> {
+        type Error = AV::Error;
+        fn traverse_address(
+            &mut self,
+            _driver: &AV::ValueDriver<'_, 'b, 'l>,
+            value: AccountAddress,
+        ) -> Result<(), Self::Error> {
+            self.0.insert(value.into());
+            Ok(())
         }
     }
-    Ok(())
+
+    MoveValue::visit_deserialize(
+        bcs_bytes,
+        fully_annotated_layout,
+        &mut UIDTraversal(&mut ids),
+    )
+    .map_err(|e| format!("Failed to deserialize. {e:?}"))?;
+    Ok(ids)
 }