[table] add information about the table types when adding new tables

Closes: diem#470
vikalpacn · Mar 17, 2023 · d5fa870 · d5fa870
1 parent d9ab076
commit d5fa870
Show file tree

Hide file tree

Showing 4 changed files with 44 additions and 9 deletions.
diff --git a/language/extensions/move-table-extension/sources/Table.move b/language/extensions/move-table-extension/sources/Table.move
@@ -18,7 +18,7 @@ module extensions::table {
     /// Create a new Table.
     public fun new<K: copy + drop, V: store>(): Table<K, V> {
         Table{
-            handle: new_table_handle(),
+            handle: new_table_handle<K, V>(),
             length: 0,
         }
     }
@@ -98,7 +98,7 @@ module extensions::table {
 
     // Primitives which take as an additional type parameter `Box<V>`, so the implementation
     // can use this to determine serialization layout.
-    native fun new_table_handle(): u128;
+    native fun new_table_handle<K, V>(): u128;
     native fun add_box<K: copy + drop, V, B>(table: &mut Table<K, V>, key: K, val: Box<V>);
     native fun borrow_box<K: copy + drop, V, B>(table: &Table<K, V>, key: K): &Box<V>;
     native fun borrow_box_mut<K: copy + drop, V, B>(table: &mut Table<K, V>, key: K): &mut Box<V>;

diff --git a/language/extensions/move-table-extension/src/lib.rs b/language/extensions/move-table-extension/src/lib.rs
@@ -11,6 +11,7 @@ use move_binary_format::errors::{PartialVMError, PartialVMResult};
 use move_core_types::{
     account_address::AccountAddress,
     gas_schedule::{GasAlgebra, GasCarrier, InternalGasUnits},
+    language_storage::TypeTag,
     value::MoveTypeLayout,
     vm_status::StatusCode,
 };
@@ -48,10 +49,31 @@ impl Display for TableHandle {
     }
 }
 
+#[derive(Clone, Debug)]
+pub struct TableInfo {
+    pub key_type: TypeTag,
+    pub value_type: TypeTag,
+}
+
+impl TableInfo {
+    pub fn new(key_type: TypeTag, value_type: TypeTag) -> Self {
+        Self {
+            key_type,
+            value_type,
+        }
+    }
+}
+
+impl Display for TableInfo {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        writeln!(f, "Table<{}, {}>", self.key_type, self.value_type)
+    }
+}
+
 /// A table change set.
 #[derive(Default)]
 pub struct TableChangeSet {
-    pub new_tables: BTreeSet<TableHandle>,
+    pub new_tables: BTreeMap<TableHandle, TableInfo>,
     pub removed_tables: BTreeSet<TableHandle>,
     pub changes: BTreeMap<TableHandle, TableChange>,
 }
@@ -115,7 +137,7 @@ const _NOT_EMPTY: u64 = (102 << 8) + _ECATEGORY_INVALID_STATE as u64;
 /// of the overall context so we can mutate while still accessing the overall context.
 #[derive(Default)]
 struct TableData {
-    new_tables: BTreeSet<TableHandle>,
+    new_tables: BTreeMap<TableHandle, TableInfo>,
     removed_tables: BTreeSet<TableHandle>,
     tables: BTreeMap<TableHandle, Table>,
 }
@@ -335,9 +357,10 @@ pub fn table_natives(table_addr: AccountAddress) -> NativeFunctionTable {
 
 fn native_new_table_handle(
     context: &mut NativeContext,
-    mut _ty_args: Vec<Type>,
+    ty_args: Vec<Type>,
     args: VecDeque<Value>,
 ) -> PartialVMResult<NativeResult> {
+    assert_eq!(ty_args.len(), 2);
     assert!(args.is_empty());
 
     let table_context = context.extensions().get::<NativeTableContext>();
@@ -351,7 +374,12 @@ fn native_new_table_handle(
     Digest::update(&mut digest, table_data.new_tables.len().to_be_bytes());
     let bytes: [u8; 16] = digest.finalize()[0..16].try_into().unwrap();
     let id = u128::from_be_bytes(bytes);
-    assert!(table_data.new_tables.insert(TableHandle(id)));
+    let key_type = context.type_to_type_tag(&ty_args[0])?;
+    let value_type = context.type_to_type_tag(&ty_args[1])?;
+    assert!(table_data
+        .new_tables
+        .insert(TableHandle(id), TableInfo::new(key_type, value_type))
+        .is_none());
 
     Ok(NativeResult::ok(
         table_context

diff --git a/language/move-vm/test-utils/src/storage.rs b/language/move-vm/test-utils/src/storage.rs
@@ -245,8 +245,12 @@ impl InMemoryStorage {
             changes,
         } = changes;
         self.tables.retain(|h, _| !removed_tables.contains(h));
-        self.tables
-            .extend(new_tables.into_iter().map(|h| (h, BTreeMap::default())));
+        self.tables.extend(
+            new_tables
+                .keys()
+                .into_iter()
+                .map(|h| (*h, BTreeMap::default())),
+        );
         for (h, c) in changes {
             assert!(
                 self.tables.contains_key(&h),

diff --git a/language/tools/move-unit-test/src/extensions.rs b/language/tools/move-unit-test/src/extensions.rs
@@ -71,7 +71,10 @@ fn print_table_extension<W: Write>(w: &mut W, extensions: &mut NativeContextExte
             writeln!(
                 w,
                 "new tables {}",
-                cs.new_tables.iter().map(|h| h.to_string()).join(", ")
+                cs.new_tables
+                    .iter()
+                    .map(|(k, v)| format!("{}<{},{}>", k, v.key_type, v.value_type))
+                    .join(", ")
             )
             .unwrap();
         }