Skip to content

Commit

Permalink
adapter(visitor): Use UIDTraversal across all execution layers
Browse files Browse the repository at this point in the history 
## 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
```
  • Loading branch information
@amnn
amnn committed Sep 30, 2024
1 parent 76cf3c6 commit 80df4a3
Show file tree
Hide file tree
Showing 4 changed files with 114 additions and 111 deletions.
12 changes: 6 additions & 6 deletions sui-execution/latest/sui-move-natives/src/object_runtime/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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)
Expand Down
71 changes: 36 additions & 35 deletions sui-execution/v0/sui-move-natives/src/object_runtime/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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::{
Expand Down Expand Up @@ -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
Expand All @@ -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)
}
71 changes: 36 additions & 35 deletions sui-execution/v1/sui-move-natives/src/object_runtime/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand All @@ -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)
}
71 changes: 36 additions & 35 deletions sui-execution/v2/sui-move-natives/src/object_runtime/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand All @@ -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)
}

0 comments on commit 80df4a3

Please sign in to comment.