fixup! fixup! fixup! fixup! Fix comments

aptos-move/aptos-intent/Cargo.toml

@@ -19,19 +19,19 @@ crate-type = ["cdylib", "rlib"]
 anyhow = { workspace = true }
 bcs = { workspace = true }
 bytes = { workspace = true }
+getrandom = { version = "0.2", features = ["js"] }
+hex = { workspace = true }
 move-binary-format = { workspace = true }
 move-bytecode-verifier = { workspace = true }
 move-core-types = { workspace = true }
+reqwest = { version = "*", features = ["blocking"] }
+reqwest-wasm = { version = "*", features = ["blocking"] }
 serde = { workspace = true }
 serde_bytes = { workspace = true }
 serde_json = { workspace = true }
-hex = { workspace = true }
-getrandom = { version = "0.2", features = ["js"] }
 wasm-bindgen = "0.2"
 wasm-bindgen-futures = "0.4.42"
-reqwest-wasm = { version = "*", features = ["blocking"] }
-reqwest = { version = "*", features = ["blocking"] }
 
 [dev-dependencies]
-e2e-move-tests = { path = "../e2e-move-tests" }
 aptos-types = { workspace = true }
+e2e-move-tests = { path = "../e2e-move-tests" }

aptos-move/aptos-intent/src/codegen.rs

@@ -6,6 +6,7 @@ use crate::{
 };
 use move_binary_format::{
     access::ScriptAccess,
+    builders::CompiledScriptBuilder,
     errors::{PartialVMError, PartialVMResult},
     file_format::*,
 };
@@ -91,16 +92,8 @@ impl Context {
     fn compile_batched_call(
         &mut self,
         call: &BatchedFunctionCall,
-        module_resolver: &BTreeMap<ModuleId, CompiledModule>,
+        func_id: FunctionHandleIndex,
     ) -> PartialVMResult<()> {
-        let target_module = match module_resolver.get(&call.module) {
-            Some(module) => module,
-            None => return Err(PartialVMError::new(StatusCode::LOOKUP_FAILED)),
-        };
-        let func_id = self
-            .script
-            .import_call_by_name(&call.function, target_module)?;
-
         let func_handle = self.script.function_handle_at(func_id).clone();
         let mut subst_mapping = BTreeMap::new();
         for (idx, ty_param) in call.ty_args.iter().enumerate() {
@@ -178,7 +171,7 @@ impl Context {
                 .map(|idx| SignatureToken::TypeParameter(*idx))
                 .collect();
 
-            let type_parameters = self.script.add_signature(Signature(inst_sig))?;
+            let type_parameters = self.add_signature(Signature(inst_sig))?;
             self.script
                 .function_instantiations
                 .push(FunctionInstantiation {
@@ -218,6 +211,25 @@ impl Context {
 
         Ok(())
     }
+
+    fn add_signature(&mut self, sig: Signature) -> PartialVMResult<SignatureIndex> {
+        Ok(SignatureIndex(match self
+            .script
+            .signatures()
+            .iter()
+            .position(|item| item == &sig)
+        {
+            Some(idx) => idx,
+            None => {
+                let idx = self.script.signatures().len();
+                if idx >= TableIndex::MAX as usize {
+                    return Err(PartialVMError::new(StatusCode::INDEX_OUT_OF_BOUNDS));
+                }
+                self.script.signatures.push(sig);
+                idx
+            },
+        } as u16))
+    }
 }
 
 #[derive(Serialize, Deserialize)]
@@ -235,21 +247,26 @@ pub fn generate_script_from_batched_calls(
     module_resolver: &BTreeMap<ModuleId, CompiledModule>,
 ) -> PartialVMResult<Vec<u8>> {
     let mut context = Context::default();
-    context.script = empty_script();
-    context.script.code.code = vec![];
-    context.script.signatures = vec![];
+    let mut script_builder = CompiledScriptBuilder::new(CompiledScript::default());
+    let call_idxs = calls
+        .iter()
+        .map(|call| {
+            let target_module = match module_resolver.get(&call.module) {
+                Some(module) => module,
+                None => return Err(PartialVMError::new(StatusCode::LOOKUP_FAILED)),
+            };
+            script_builder.import_call_by_name(&call.function, target_module)
+        })
+        .collect::<PartialVMResult<Vec<_>>>()?;
+    context.script = script_builder.into_inner();
     context.add_signers(signer_count)?;
     context.allocate_parameters(calls)?;
-    for call in calls.iter() {
-        context.compile_batched_call(call, module_resolver)?;
+    for (call, fh_idx) in calls.iter().zip(call_idxs.into_iter()) {
+        context.compile_batched_call(call, fh_idx)?;
     }
     context.script.code.code.push(Bytecode::Ret);
-    context.script.parameters = context
-        .script
-        .add_signature(Signature(context.parameters.clone()))?;
-    context.script.code.locals = context
-        .script
-        .add_signature(Signature(context.locals.clone()))?;
+    context.script.parameters = context.add_signature(Signature(context.parameters.clone()))?;
+    context.script.code.locals = context.add_signature(Signature(context.locals.clone()))?;
     move_bytecode_verifier::verify_script(&context.script).map_err(|err| {
         err.to_partial()
             .with_message(format!("{:?}", context.script))

third_party/move/move-binary-format/src/builders.rs

@@ -2,30 +2,111 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use crate::{
-    access::ModuleAccess,
+    access::{ModuleAccess, ScriptAccess},
     errors::{PartialVMError, PartialVMResult},
     file_format::*,
 };
 use move_core_types::{
-    account_address::AccountAddress, identifier::IdentStr, language_storage::ModuleId,
+    account_address::AccountAddress,
+    identifier::{IdentStr, Identifier},
+    language_storage::ModuleId,
     vm_status::StatusCode,
 };
+use std::collections::{btree_map::Entry, BTreeMap};
 
-fn get_or_add<T: PartialEq>(pool: &mut Vec<T>, val: T) -> PartialVMResult<usize> {
-    match pool.iter().position(|elem| elem == &val) {
-        Some(idx) => Ok(idx),
-        None => {
+pub struct CompiledScriptBuilder {
+    script: CompiledScript,
+    address_pool: BTreeMap<AccountAddress, usize>,
+    identifier_pool: BTreeMap<Identifier, usize>,
+    module_pool: BTreeMap<(AddressIdentifierIndex, IdentifierIndex), usize>,
+    function_pool: BTreeMap<(ModuleHandleIndex, IdentifierIndex), usize>,
+    struct_pool: BTreeMap<(ModuleHandleIndex, IdentifierIndex), usize>,
+    signature_pool: BTreeMap<Signature, usize>,
+}
+
+fn get_or_add_impl<T, I, F>(
+    pool: &mut Vec<T>,
+    idx_map: &mut BTreeMap<I, usize>,
+    val: F,
+    idx: I,
+) -> PartialVMResult<usize>
+where
+    I: Ord,
+    F: FnOnce() -> T,
+{
+    match idx_map.entry(idx) {
+        Entry::Occupied(i) => Ok(*i.get()),
+        Entry::Vacant(entry) => {
             let idx = pool.len();
             if pool.len() >= TableIndex::MAX as usize {
                 return Err(PartialVMError::new(StatusCode::INDEX_OUT_OF_BOUNDS));
             }
-            pool.push(val);
+            pool.push(val());
+            entry.insert(idx);
             Ok(idx)
         },
     }
 }
 
-impl CompiledScript {
+fn get_or_add<T: Ord + Clone>(
+    pool: &mut Vec<T>,
+    idx_map: &mut BTreeMap<T, usize>,
+    val: T,
+) -> PartialVMResult<usize> {
+    let val_cloned = val.clone();
+    get_or_add_impl(pool, idx_map, || val, val_cloned)
+}
+
+impl CompiledScriptBuilder {
+    pub fn new(script: CompiledScript) -> Self {
+        let address_pool = script
+            .address_identifiers()
+            .iter()
+            .enumerate()
+            .map(|(idx, addr)| (*addr, idx))
+            .collect();
+        let identifier_pool = script
+            .identifiers()
+            .iter()
+            .enumerate()
+            .map(|(idx, ident)| (ident.to_owned(), idx))
+            .collect();
+        let module_pool = script
+            .module_handles()
+            .iter()
+            .enumerate()
+            .map(|(idx, handle)| ((handle.address, handle.name), idx))
+            .collect();
+        let function_pool = script
+            .function_handles()
+            .iter()
+            .enumerate()
+            .map(|(idx, handle)| ((handle.module, handle.name), idx))
+            .collect();
+        let struct_pool = script
+            .struct_handles()
+            .iter()
+            .enumerate()
+            .map(|(idx, handle)| ((handle.module, handle.name), idx))
+            .collect();
+        let signature_pool = script
+            .signatures()
+            .iter()
+            .enumerate()
+            .map(|(idx, signature)| (signature.clone(), idx))
+            .collect();
+
+        Self {
+            script,
+            address_pool,
+            identifier_pool,
+            module_pool,
+            function_pool,
+            struct_pool,
+            signature_pool,
+        }
+    }
+
     pub fn import_address(
         &mut self,
         module: &CompiledModule,
@@ -38,15 +119,24 @@ impl CompiledScript {
         &mut self,
         address: &AccountAddress,
     ) -> PartialVMResult<AddressIdentifierIndex> {
-        get_or_add(&mut self.address_identifiers, *address)
-            .map(|idx| AddressIdentifierIndex(idx as u16))
+        get_or_add(
+            &mut self.script.address_identifiers,
+            &mut self.address_pool,
+            *address,
+        )
+        .map(|idx| AddressIdentifierIndex(idx as u16))
     }
 
     pub fn import_identifier_by_name(
         &mut self,
         ident: &IdentStr,
     ) -> PartialVMResult<IdentifierIndex> {
-        get_or_add(&mut self.identifiers, ident.to_owned()).map(|idx| IdentifierIndex(idx as u16))
+        get_or_add(
+            &mut self.script.identifiers,
+            &mut self.identifier_pool,
+            ident.to_owned(),
+        )
+        .map(|idx| IdentifierIndex(idx as u16))
     }
 
     pub fn import_identifier(
@@ -79,8 +169,13 @@ impl CompiledScript {
         address: AddressIdentifierIndex,
         name: IdentifierIndex,
     ) -> PartialVMResult<ModuleHandleIndex> {
-        get_or_add(&mut self.module_handles, ModuleHandle { address, name })
-            .map(|idx| ModuleHandleIndex(idx as u16))
+        get_or_add_impl(
+            &mut self.script.module_handles,
+            &mut self.module_pool,
+            || ModuleHandle { address, name },
+            (address, name),
+        )
+        .map(|idx| ModuleHandleIndex(idx as u16))
     }
 
     pub fn import_struct(
@@ -91,26 +186,17 @@ impl CompiledScript {
         let handle = module.struct_handle_at(idx);
         let module_id = self.import_module(module, handle.module)?;
         let name = self.import_identifier(module, handle.name)?;
-        let idx = match self
-            .struct_handles
-            .iter()
-            .position(|elem| elem.module == module_id && elem.name == name)
-        {
-            Some(idx) => Ok(idx),
-            None => {
-                let idx = self.struct_handles.len();
-                if self.struct_handles.len() >= TableIndex::MAX as usize {
-                    return Err(PartialVMError::new(StatusCode::INDEX_OUT_OF_BOUNDS));
-                }
-                self.struct_handles.push(StructHandle {
-                    module: module_id,
-                    name,
-                    abilities: handle.abilities,
-                    type_parameters: handle.type_parameters.clone(),
-                });
-                Ok(idx)
+        let idx = get_or_add_impl(
+            &mut self.script.struct_handles,
+            &mut self.struct_pool,
+            || StructHandle {
+                module: module_id,
+                name,
+                abilities: handle.abilities,
+                type_parameters: handle.type_parameters.clone(),
             },
-        }?;
+            (module_id, name),
+        )?;
         Ok(StructHandleIndex(idx as u16))
     }
 
@@ -164,7 +250,12 @@ impl CompiledScript {
     }
 
     pub fn add_signature(&mut self, signature: Signature) -> PartialVMResult<SignatureIndex> {
-        get_or_add(&mut self.signatures, signature).map(|idx| SignatureIndex(idx as u16))
+        get_or_add(
+            &mut self.script.signatures,
+            &mut self.signature_pool,
+            signature,
+        )
+        .map(|idx| SignatureIndex(idx as u16))
     }
 
     pub fn import_call(
@@ -175,28 +266,25 @@ impl CompiledScript {
         let handle = module.function_handle_at(idx);
         let module_id = self.import_module(module, handle.module)?;
         let name = self.import_identifier(module, handle.name)?;
-        let idx = match self
-            .function_handles
-            .iter()
-            .position(|elem| elem.module == module_id && elem.name == name)
-        {
-            Some(idx) => Ok(idx),
+        let idx = match self.function_pool.get(&(module_id, name)) {
+            Some(idx) => Ok(*idx),
             None => {
-                let idx = self.struct_handles.len();
-                if self.function_handles.len() >= TableIndex::MAX as usize {
+                let idx = self.script.struct_handles.len();
+                if self.script.function_handles.len() >= TableIndex::MAX as usize {
                     return Err(PartialVMError::new(StatusCode::INDEX_OUT_OF_BOUNDS));
                 }
                 let parameters = self.import_signatures(module, handle.parameters)?;
                 let return_ = self.import_signatures(module, handle.return_)?;
 
-                self.function_handles.push(FunctionHandle {
+                self.script.function_handles.push(FunctionHandle {
                     module: module_id,
                     name,
                     parameters,
                     return_,
                     type_parameters: handle.type_parameters.clone(),
                     access_specifiers: handle.access_specifiers.clone(),
                 });
+                self.function_pool.insert((module_id, name), idx);
                 Ok(idx)
             },
         }?;
@@ -215,4 +303,8 @@ impl CompiledScript {
         }
         Err(PartialVMError::new(StatusCode::LOOKUP_FAILED))
     }
+
+    pub fn into_inner(self) -> CompiledScript {
+        self.script
+    }
 }

third_party/move/move-binary-format/src/lib.rs

@@ -14,7 +14,7 @@ pub mod compatibility;
 pub mod compatibility_legacy;
 #[macro_use]
 pub mod errors;
-mod builders;
+pub mod builders;
 pub mod check_complexity;
 pub mod constant;
 pub mod control_flow_graph;