Take into account table changes while making the genesis change set

Closes: aptos-labs#804

aptos-move/aptos-vm/src/move_vm_ext/session.rs

@@ -20,8 +20,8 @@ use move_core_types::{
     effects::{ChangeSet as MoveChangeSet, Event as MoveEvent},
     vm_status::{StatusCode, VMStatus},
 };
-use move_table_extension::NativeTableContext;
-use move_vm_runtime::{native_extensions::NativeContextExtensions, session::Session};
+use move_table_extension::{NativeTableContext, TableChange, TableChangeSet};
+use move_vm_runtime::session::Session;
 use serde::{Deserialize, Serialize};
 use std::{
     convert::TryInto,
@@ -96,12 +96,17 @@ where
         Self { inner }
     }
 
-    pub fn finish(self) -> VMResult<SessionOutput<'r>> {
-        let (change_set, events, extensions) = self.inner.finish_with_extensions()?;
+    pub fn finish(self) -> VMResult<SessionOutput> {
+        let (change_set, events, mut extensions) = self.inner.finish_with_extensions()?;
+        let table_context: NativeTableContext = extensions.remove();
+        let table_change_set = table_context
+            .into_change_set()
+            .map_err(|e| e.finish(Location::Undefined))?;
+
         Ok(SessionOutput {
             change_set,
             events,
-            extensions,
+            table_change_set,
         })
     }
 }
@@ -120,15 +125,15 @@ impl<'r, 'l, S> DerefMut for SessionExt<'r, 'l, S> {
     }
 }
 
-pub struct SessionOutput<'r> {
+pub struct SessionOutput {
     pub change_set: MoveChangeSet,
     pub events: Vec<MoveEvent>,
-    pub extensions: NativeContextExtensions<'r>,
+    pub table_change_set: TableChangeSet,
 }
 
-impl<'r> SessionOutput<'r> {
+impl SessionOutput {
     pub fn into_change_set<C: AccessPathCache>(
-        mut self,
+        self,
         ap_cache: &mut C,
     ) -> Result<ChangeSet, VMStatus> {
         let mut out_write_set = WriteSetMut::new(Vec::new());
@@ -141,19 +146,48 @@ impl<'r> SessionOutput<'r> {
             &mut out_events,
         )?;
 
-        let table_context: NativeTableContext = self.extensions.remove();
-        let table_changeset = table_context
-            .into_change_set()
-            .map_err(|e| e.finish(Location::Undefined))?;
-        convert_table_changeset(table_changeset, &mut out_write_set)?;
+        convert_table_changeset(self.table_change_set, &mut out_write_set)?;
 
         let ws = out_write_set
             .freeze()
             .map_err(|_| VMStatus::Error(StatusCode::DATA_FORMAT_ERROR))?;
         Ok(ChangeSet::new(ws, out_events))
     }
 
-    pub fn unpack(self) -> (MoveChangeSet, Vec<MoveEvent>, NativeContextExtensions<'r>) {
-        (self.change_set, self.events, self.extensions)
+    pub fn unpack(self) -> (MoveChangeSet, Vec<MoveEvent>, TableChangeSet) {
+        (self.change_set, self.events, self.table_change_set)
+    }
+
+    pub fn squash(&mut self, other: Self) -> Result<(), VMStatus> {
+        self.change_set
+            .squash(other.change_set)
+            .map_err(|_| VMStatus::Error(StatusCode::DATA_FORMAT_ERROR))?;
+        self.events.extend(other.events.into_iter());
+
+        // Squash the table changes
+        self.table_change_set
+            .new_tables
+            .extend(other.table_change_set.new_tables);
+        for removed_table in &self.table_change_set.removed_tables {
+            self.table_change_set.new_tables.remove(removed_table);
+        }
+        // There's chance that a table is added in `self`, and an item is added to that table in
+        // `self`, and later the item is deleted in `other`, netting to a NOOP for that item,
+        // but this is an tricky edge case that we don't expect to happen too much, it doesn't hurt
+        // too much to just keep the deletion. It's safe as long as we do it that way consistently.
+        self.table_change_set
+            .removed_tables
+            .extend(other.table_change_set.removed_tables.into_iter());
+        for (handle, changes) in other.table_change_set.changes.into_iter() {
+            let my_changes = self
+                .table_change_set
+                .changes
+                .entry(handle)
+                .or_insert(TableChange {
+                    entries: Default::default(),
+                });
+            my_changes.entries.extend(changes.entries.into_iter());
+        }
+        Ok(())
     }
 }

aptos-move/vm-genesis/src/lib.rs

@@ -21,7 +21,6 @@ use aptos_types::{
     transaction::{authenticator::AuthenticationKey, ChangeSet, Transaction, WriteSetPayload},
 };
 use aptos_vm::{
-    convert_changeset_and_events,
     data_cache::{IntoMoveResolver, StateViewCache},
     move_vm_ext::{MoveVmExt, SessionExt, SessionId},
 };
@@ -110,7 +109,7 @@ pub fn encode_genesis_change_set(
     create_and_initialize_validators(&mut session, validators);
     reconfigure(&mut session);
 
-    let (mut changeset1, mut events1, _) = session.finish().unwrap().unpack();
+    let mut session1_out = session.finish().unwrap();
 
     let state_view = GenesisStateView::new();
     let data_cache = StateViewCache::new(&state_view).into_move_resolver();
@@ -122,16 +121,17 @@ pub fn encode_genesis_change_set(
     let mut session = move_vm.new_session(&data_cache, SessionId::genesis(id2));
 
     publish_stdlib(&mut session, Modules::new(stdlib_modules.iter()));
-    let (changeset2, events2, _) = session.finish().unwrap().unpack();
+    let session2_out = session.finish().unwrap();
 
-    changeset1.squash(changeset2).unwrap();
-    events1.extend(events2);
+    session1_out.squash(session2_out).unwrap();
+    let change_set = session1_out.into_change_set(&mut ()).unwrap();
 
-    let (write_set, events) = convert_changeset_and_events(changeset1, events1).unwrap();
-
-    assert!(!write_set.iter().any(|(_, op)| op.is_deletion()));
-    verify_genesis_write_set(&events);
-    ChangeSet::new(write_set, events)
+    assert!(!change_set
+        .write_set()
+        .iter()
+        .any(|(_, op)| op.is_deletion()));
+    verify_genesis_write_set(change_set.events());
+    change_set
 }
 
 fn exec_function(