Add ObjectRootAncestorMap (MystenLabs#2471)

crates/sui-adapter/src/adapter.rs

@@ -41,6 +41,7 @@ use std::{
     fmt::Debug,
 };
 
+use crate::object_root_ancestor_map::ObjectRootAncestorMap;
 pub use move_vm_runtime::move_vm::MoveVM;
 
 pub fn new_move_vm(natives: NativeFunctionTable) -> Result<MoveVM, SuiError> {
@@ -869,61 +870,25 @@ fn check_child_object_of_shared_object(
     object_type_map: &BTreeMap<ObjectID, ModuleId>,
     current_module: ModuleId,
 ) -> SuiResult {
-    // ancestor_map is a cache that remembers the top ancestor of each object.
-    // Top ancestor is the root object at the top in the object ownership chain whose
-    // parent is no longer an object.
-    let mut ancestor_map: BTreeMap<ObjectID, ObjectID> = BTreeMap::new();
-    for (object_id, obj) in objects {
-        // We only need to look at objects that are owned by other objects.
-        if !matches!(obj.borrow().owner, Owner::ObjectOwner(_)) {
+    let object_owner_map = objects
+        .iter()
+        .map(|(id, obj)| (*id, obj.borrow().owner))
+        .collect();
+    let ancestor_map = ObjectRootAncestorMap::new(&object_owner_map)?;
+    for (object_id, owner) in object_owner_map {
+        // We are only interested in objects owned by objects.
+        if !matches!(owner, Owner::ObjectOwner(..)) {
             continue;
         }
-        // We trace the object up through the ownership chain, until we either hit
-        // the top (no more parent object), or we hit an object that's already
-        // in the `ancestor_map` (because we visited this object previously).
-        // We then have a pair between this object and its top ancestor. If the top
-        // ancestor is a shared object, we check on their types.
-        let mut ancestor_stack = vec![];
-        let mut cur_id = *object_id;
-        let mut cur_obj = obj;
-        let ancestor_id = loop {
-            if let Some(ancestor) = ancestor_map.get(&cur_id) {
-                break *ancestor;
-            }
-            ancestor_stack.push(cur_id);
-            match cur_obj.borrow().owner {
-                Owner::ObjectOwner(parent_id) => {
-                    cur_obj =
-                        objects
-                            .get(&parent_id.into())
-                            .ok_or(SuiError::MissingObjectOwner {
-                                child_id: cur_id,
-                                parent_id: parent_id.into(),
-                            })?;
-                    cur_id = parent_id.into();
-                    fp_ensure!(
-                        cur_id != ancestor_stack[0],
-                        SuiError::CircularObjectOwnership
-                    );
-                }
-                Owner::AddressOwner(_) | Owner::Immutable | Owner::Shared => {
-                    break cur_id;
-                }
-            };
-        };
-        // For each ancestor we have visited, cache their top ancestor so that if we
-        // ever visit them in the future, we know the answer.
-        while let Some(id) = ancestor_stack.pop() {
-            ancestor_map.insert(id, ancestor_id);
-        }
-        // Check the orphan rule.
-        if objects.get(&ancestor_id).unwrap().borrow().is_shared() {
-            let child_module = object_type_map.get(object_id).unwrap();
+        let (ancestor_id, ancestor_owner) = ancestor_map.get_root_ancestor(&object_id)?;
+        if ancestor_owner.is_shared() {
+            // unwrap safe because the object ID exists in object_owner_map.
+            let child_module = object_type_map.get(&object_id).unwrap();
             let ancestor_module = object_type_map.get(&ancestor_id).unwrap();
             fp_ensure!(
                 child_module == &current_module || ancestor_module == &current_module,
                 SuiError::InvalidSharedChildUse {
-                    child: *object_id,
+                    child: object_id,
                     child_module: current_module.to_string(),
                     ancestor: ancestor_id,
                     ancestor_module: ancestor_module.to_string(),

crates/sui-adapter/src/lib.rs

@@ -4,3 +4,4 @@
 pub mod adapter;
 pub mod bytecode_rewriter;
 pub mod genesis;
+pub mod object_root_ancestor_map;

crates/sui-adapter/src/object_root_ancestor_map.rs

@@ -0,0 +1,67 @@
+// Copyright (c) 2022, Mysten Labs, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+use std::collections::BTreeMap;
+use sui_types::base_types::ObjectID;
+use sui_types::error::{SuiError, SuiResult};
+use sui_types::fp_ensure;
+use sui_types::object::Owner;
+
+/// A structure that maps from object ID to its root ancestor object.
+/// For an object A that's owned by another object B, the root ancestor object of A is the root
+/// ancestor object of B;
+/// For an object that's not owned by another object, its root ancestor is itself.
+/// Given a list map of object ID and the owner of the object, this data structure builds up
+/// the ancestor map efficiently (O(N)). For each object ID, `root_ancestor_map` stores the object's
+/// root ancestor object's ID and owner.
+pub struct ObjectRootAncestorMap {
+    root_ancestor_map: BTreeMap<ObjectID, (ObjectID, Owner)>,
+}
+
+impl ObjectRootAncestorMap {
+    pub fn new(direct_owner_map: &BTreeMap<ObjectID, Owner>) -> SuiResult<Self> {
+        let mut root_ancestor_map = BTreeMap::new();
+        for (id, owner) in direct_owner_map {
+            // All the objects we will visit while walking up the ancestor chain.
+            // Remember them so that we could set each of their root ancestor to the same result.
+            let mut stack = vec![];
+
+            let mut cur_id = *id;
+            let mut cur_owner = *owner;
+            let ancestor_info = loop {
+                if let Some(ancestor) = root_ancestor_map.get(&cur_id) {
+                    break *ancestor;
+                }
+                stack.push(cur_id);
+                match cur_owner {
+                    Owner::ObjectOwner(parent_id) => {
+                        cur_owner = *direct_owner_map.get(&parent_id.into()).ok_or(
+                            SuiError::MissingObjectOwner {
+                                child_id: cur_id,
+                                parent_id: parent_id.into(),
+                            },
+                        )?;
+                        cur_id = parent_id.into();
+                        fp_ensure!(cur_id != stack[0], SuiError::CircularObjectOwnership);
+                    }
+                    Owner::AddressOwner(_) | Owner::Immutable | Owner::Shared => {
+                        break (cur_id, cur_owner);
+                    }
+                };
+            };
+            while let Some(object_id) = stack.pop() {
+                root_ancestor_map.insert(object_id, ancestor_info);
+            }
+        }
+        Ok(Self { root_ancestor_map })
+    }
+
+    pub fn get_root_ancestor(&self, object_id: &ObjectID) -> SuiResult<(ObjectID, Owner)> {
+        Ok(*self
+            .root_ancestor_map
+            .get(object_id)
+            .ok_or(SuiError::ObjectNotFound {
+                object_id: *object_id,
+            })?)
+    }
+}