install table extension

api/types/src/convert.rs

@@ -208,15 +208,9 @@ impl<'a, R: MoveResolverExt + ?Sized> MoveConverter<'a, R> {
             StateKey::AccessPath(access_path) => {
                 self.try_access_path_into_write_set_change(access_path, op)
             }
-            // We should not expect account address here.
-            StateKey::AccountAddressKey(_) => Err(format_err!(
-                "Can't convert account address key {:?} to WriteSetChange",
-                state_key
-            )),
-            StateKey::Raw(_) => Err(format_err!(
-                "Can't convert account raw key {:?} to WriteSetChange",
-                state_key
-            )),
+            StateKey::AccountAddressKey(_) | StateKey::Raw(_) | StateKey::TableItem { .. } => Err(
+                format_err!("Can't convert state key {:?} to WriteSetChange", state_key),
+            ),
         }
     }
 

aptos-move/aptos-vm/src/aptos_vm_impl.rs

@@ -40,6 +40,7 @@ use move_core_types::{
     resolver::ResourceResolver,
     value::{serialize_values, MoveValue},
 };
+use move_table_extension::TableChangeSet;
 use move_vm_runtime::{logging::expect_no_verification_errors, session::Session};
 use move_vm_types::gas_schedule::{calculate_intrinsic_gas, GasStatus};
 use std::{convert::TryFrom, sync::Arc};
@@ -549,17 +550,28 @@ pub fn convert_changeset_and_events(
     changeset: MoveChangeSet,
     events: Vec<MoveEvent>,
 ) -> Result<(WriteSet, Vec<ContractEvent>), VMStatus> {
-    convert_changeset_and_events_cached(&mut (), changeset, events)
+    let mut out_write_set = WriteSetMut::new(vec![]);
+    let mut out_events = Vec::new();
+    convert_changeset_and_events_cached(
+        &mut (),
+        changeset,
+        events,
+        &mut out_write_set,
+        &mut out_events,
+    )?;
+    let ws = out_write_set
+        .freeze()
+        .map_err(|_| VMStatus::Error(StatusCode::DATA_FORMAT_ERROR))?;
+    Ok((ws, out_events))
 }
 
 pub fn convert_changeset_and_events_cached<C: AccessPathCache>(
     ap_cache: &mut C,
     changeset: MoveChangeSet,
     events: Vec<MoveEvent>,
-) -> Result<(WriteSet, Vec<ContractEvent>), VMStatus> {
-    // TODO: Cache access path computations if necessary.
-    let mut ops = vec![];
-
+    out_write_set: &mut WriteSetMut,
+    out_events: &mut Vec<ContractEvent>,
+) -> Result<(), VMStatus> {
     for (addr, account_changeset) in changeset.into_inner() {
         let (modules, resources) = account_changeset.into_inner();
         for (struct_tag, blob_opt) in resources {
@@ -568,7 +580,7 @@ pub fn convert_changeset_and_events_cached<C: AccessPathCache>(
                 None => WriteOp::Deletion,
                 Some(blob) => WriteOp::Value(blob),
             };
-            ops.push((StateKey::AccessPath(ap), op))
+            out_write_set.push((StateKey::AccessPath(ap), op))
         }
 
         for (name, blob_opt) in modules {
@@ -578,14 +590,10 @@ pub fn convert_changeset_and_events_cached<C: AccessPathCache>(
                 Some(blob) => WriteOp::Value(blob),
             };
 
-            ops.push((StateKey::AccessPath(ap), op))
+            out_write_set.push((StateKey::AccessPath(ap), op))
         }
     }
 
-    let ws = WriteSetMut::new(ops)
-        .freeze()
-        .map_err(|_| VMStatus::Error(StatusCode::DATA_FORMAT_ERROR))?;
-
     let events = events
         .into_iter()
         .map(|(guid, seq_num, ty_tag, blob)| {
@@ -595,7 +603,26 @@ pub fn convert_changeset_and_events_cached<C: AccessPathCache>(
         })
         .collect::<Result<Vec<_>, VMStatus>>()?;
 
-    Ok((ws, events))
+    out_events.extend(events.into_iter());
+    Ok(())
+}
+
+pub fn convert_table_changeset(
+    table_changeset: TableChangeSet,
+    out_write_set: &mut WriteSetMut,
+) -> Result<(), VMStatus> {
+    for (handle, change) in table_changeset.changes {
+        for (key, value_opt) in change.entries {
+            let state_key = StateKey::table_item(handle.0, key);
+            if let Some(bytes) = value_opt {
+                out_write_set.push((state_key, WriteOp::Value(bytes)))
+            } else {
+                out_write_set.push((state_key, WriteOp::Deletion))
+            }
+        }
+    }
+
+    Ok(())
 }
 
 pub(crate) fn charge_global_write_gas_usage<R: MoveResolverExt>(

aptos-move/aptos-vm/src/data_cache.rs

@@ -19,7 +19,7 @@ use fail::fail_point;
 use move_binary_format::errors::*;
 use move_core_types::{
     account_address::AccountAddress,
-    gas_schedule::{GasCarrier, InternalGasUnits},
+    gas_schedule::{GasAlgebra, GasCarrier, InternalGasUnits},
     language_storage::{ModuleId, StructTag},
     resolver::{ModuleResolver, ResourceResolver},
 };
@@ -155,10 +155,10 @@ impl<'a, S: StateView> ResourceResolver for RemoteStorage<'a, S> {
 impl<'a, S: StateView> TableResolver for RemoteStorage<'a, S> {
     fn resolve_table_entry(
         &self,
-        _handle: &TableHandle,
-        _key: &[u8],
+        handle: &TableHandle,
+        key: &[u8],
     ) -> Result<Option<Vec<u8>>, Error> {
-        todo!()
+        self.get_state_value(&StateKey::table_item(handle.0, key.to_vec()))
     }
 
     fn operation_cost(
@@ -167,7 +167,7 @@ impl<'a, S: StateView> TableResolver for RemoteStorage<'a, S> {
         _key_size: usize,
         _val_size: usize,
     ) -> InternalGasUnits<GasCarrier> {
-        todo!()
+        InternalGasUnits::new(1)
     }
 }
 

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

@@ -2,21 +2,25 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use crate::{
-    access_path_cache::AccessPathCache, aptos_vm_impl::convert_changeset_and_events_cached,
-    move_vm_ext::MoveResolverExt, transaction_metadata::TransactionMetadata,
+    access_path_cache::AccessPathCache,
+    aptos_vm_impl::{convert_changeset_and_events_cached, convert_table_changeset},
+    move_vm_ext::MoveResolverExt,
+    transaction_metadata::TransactionMetadata,
 };
 use aptos_crypto::{hash::CryptoHash, HashValue};
 use aptos_crypto_derive::{BCSCryptoHash, CryptoHasher};
 use aptos_types::{
     block_metadata::BlockMetadata,
     transaction::{ChangeSet, SignatureCheckedTransaction},
+    write_set::WriteSetMut,
 };
-use move_binary_format::errors::VMResult;
+use move_binary_format::errors::{Location, VMResult};
 use move_core_types::{
     account_address::AccountAddress,
     effects::{ChangeSet as MoveChangeSet, Event as MoveEvent},
-    vm_status::VMStatus,
+    vm_status::{StatusCode, VMStatus},
 };
+use move_table_extension::NativeTableContext;
 use move_vm_runtime::{native_extensions::NativeContextExtensions, session::Session};
 use serde::{Deserialize, Serialize};
 use std::{
@@ -124,12 +128,29 @@ pub struct SessionOutput<'r> {
 
 impl<'r> SessionOutput<'r> {
     pub fn into_change_set<C: AccessPathCache>(
-        self,
+        mut self,
         ap_cache: &mut C,
     ) -> Result<ChangeSet, VMStatus> {
-        // TODO: consider table change set from the table extension
-        convert_changeset_and_events_cached(ap_cache, self.change_set, self.events)
-            .map(|(write_set, events)| ChangeSet::new(write_set, events))
+        let mut out_write_set = WriteSetMut::new(Vec::new());
+        let mut out_events = Vec::new();
+        convert_changeset_and_events_cached(
+            ap_cache,
+            self.change_set,
+            self.events,
+            &mut out_write_set,
+            &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)?;
+
+        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>) {

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

@@ -6,6 +6,7 @@ use crate::{
     natives::aptos_natives,
 };
 use move_binary_format::errors::VMResult;
+use move_table_extension::NativeTableContext;
 use move_vm_runtime::{move_vm::MoveVM, native_extensions::NativeContextExtensions};
 use std::ops::Deref;
 
@@ -23,10 +24,10 @@ impl MoveVmExt {
     pub fn new_session<'r, S: MoveResolverExt>(
         &self,
         remote: &'r S,
-        _session_id: SessionId,
+        session_id: SessionId,
     ) -> SessionExt<'r, '_, S> {
-        // TODO: install table extension
-        let extensions = NativeContextExtensions::<'r>::default();
+        let mut extensions = NativeContextExtensions::default();
+        extensions.add(NativeTableContext::new(session_id.as_uuid(), remote));
 
         SessionExt::new(self.inner.new_session_with_extensions(remote, extensions))
     }

aptos-move/aptos-vm/src/natives.rs

@@ -8,5 +8,6 @@ pub fn aptos_natives() -> NativeFunctionTable {
     move_stdlib::natives::all_natives(CORE_CODE_ADDRESS)
         .into_iter()
         .chain(framework::natives::all_natives(CORE_CODE_ADDRESS))
+        .chain(move_table_extension::table_natives(CORE_CODE_ADDRESS))
         .collect()
 }

execution/executor/src/components/apply_chunk_output.rs

@@ -402,7 +402,7 @@ pub fn process_write_set(
             }
             // For now, we only support write set with access path, this needs to be updated once
             // we support table items
-            StateKey::Raw(_) => process_state_key_write_op(
+            StateKey::Raw(_) | StateKey::TableItem { .. } => process_state_key_write_op(
                 transaction,
                 state_cache,
                 &mut updated_keys,

testsuite/generate-format/tests/staged/aptos.yaml

@@ -136,6 +136,11 @@ StateKey:
         NEWTYPE:
           TYPENAME: AccessPath
     2:
+      TableItem:
+        STRUCT:
+          - handle: U128
+          - key: BYTES
+    3:
       Raw:
         NEWTYPE: BYTES
 StructTag:

testsuite/generate-format/tests/staged/consensus.yaml

@@ -303,6 +303,11 @@ StateKey:
         NEWTYPE:
           TYPENAME: AccessPath
     2:
+      TableItem:
+        STRUCT:
+          - handle: U128
+          - key: BYTES
+    3:
       Raw:
         NEWTYPE: BYTES
 StructTag:

types/src/state_store/state_key.rs

@@ -11,6 +11,7 @@ use move_core_types::account_address::AccountAddress;
 use num_derive::{FromPrimitive, ToPrimitive};
 use num_traits::FromPrimitive;
 use serde::{Deserialize, Serialize};
+use std::convert::TryInto;
 use thiserror::Error;
 
 #[derive(
@@ -20,6 +21,11 @@ use thiserror::Error;
 pub enum StateKey {
     AccountAddressKey(AccountAddress),
     AccessPath(AccessPath),
+    TableItem {
+        handle: u128,
+        #[serde(with = "serde_bytes")]
+        key: Vec<u8>,
+    },
     // Only used for testing
     #[serde(with = "serde_bytes")]
     Raw(Vec<u8>),
@@ -30,6 +36,7 @@ pub enum StateKey {
 pub enum StateKeyTag {
     AccountAddress,
     AccessPath,
+    TableItem,
     Raw = 255,
 }
 
@@ -45,6 +52,11 @@ impl StateKey {
             StateKey::AccessPath(access_path) => {
                 (StateKeyTag::AccessPath, bcs::to_bytes(access_path)?)
             }
+            StateKey::TableItem { handle, key } => {
+                let mut bytes = handle.to_be_bytes().to_vec();
+                bytes.extend(key);
+                (StateKeyTag::TableItem, bytes)
+            }
             StateKey::Raw(raw_bytes) => (StateKeyTag::Raw, raw_bytes.to_vec()),
         };
         out.push(prefix as u8);
@@ -65,9 +77,23 @@ impl StateKey {
                 Ok(StateKey::AccountAddressKey(bcs::from_bytes(&val[1..])?))
             }
             StateKeyTag::AccessPath => Ok(StateKey::AccessPath(bcs::from_bytes(&val[1..])?)),
+            StateKeyTag::TableItem => {
+                const HANDLE_SIZE: usize = std::mem::size_of::<u128>();
+                let handle = u128::from_be_bytes(
+                    val[1..1 + HANDLE_SIZE]
+                        .try_into()
+                        .expect("Bytes too short."),
+                );
+                let key = val[1 + HANDLE_SIZE..].to_vec();
+                Ok(StateKey::table_item(handle, key))
+            }
             StateKeyTag::Raw => Ok(StateKey::Raw(val[1..].to_vec())),
         }
     }
+
+    pub fn table_item(handle: u128, key: Vec<u8>) -> Self {
+        StateKey::TableItem { handle, key }
+    }
 }
 
 impl CryptoHash for StateKey {