From 86580da4b02e7d4399e84976f8469f14540604a4 Mon Sep 17 00:00:00 2001
From: runtianz <runtian@aptoslabs.com>
Date: Fri, 25 Oct 2024 12:49:31 -0700
Subject: [PATCH] add arg

---
 crates/transaction-generator-lib/src/args.rs  |  6 +++++
 .../src/publishing/module_simple.rs           | 22 +++++++++----------
 2 files changed, 17 insertions(+), 11 deletions(-)

diff --git a/crates/transaction-generator-lib/src/args.rs b/crates/transaction-generator-lib/src/args.rs
index fb0cc7779d427..904fb1b1eeb75 100644
--- a/crates/transaction-generator-lib/src/args.rs
+++ b/crates/transaction-generator-lib/src/args.rs
@@ -74,6 +74,8 @@ pub enum TransactionTypeArg {
     SmartTablePicture1MWith1KChangeExceedsLimit,
     DeserializeU256,
     SimpleScript,
+    PermissionedTransfer,
+    APTTransfer,
 }
 
 impl TransactionTypeArg {
@@ -326,6 +328,10 @@ impl TransactionTypeArg {
             },
             TransactionTypeArg::DeserializeU256 => call_custom_module(EntryPoints::DeserializeU256),
             TransactionTypeArg::SimpleScript => call_custom_module(EntryPoints::SimpleScript),
+            TransactionTypeArg::PermissionedTransfer => {
+                call_custom_module(EntryPoints::APTPermissionedTransfer)
+            },
+            TransactionTypeArg::APTTransfer => call_custom_module(EntryPoints::APTTransfer),
         }
     }
 
diff --git a/crates/transaction-generator-lib/src/publishing/module_simple.rs b/crates/transaction-generator-lib/src/publishing/module_simple.rs
index f52e3d9f1afa8..a8875bb1fa81b 100644
--- a/crates/transaction-generator-lib/src/publishing/module_simple.rs
+++ b/crates/transaction-generator-lib/src/publishing/module_simple.rs
@@ -396,9 +396,10 @@ impl EntryPoints {
             },
             EntryPoints::IncGlobalMilestoneAggV2 { .. }
             | EntryPoints::CreateGlobalMilestoneAggV2 { .. } => "counter_with_milestone",
-            | EntryPoints::DeserializeU256 => "bcs_stream",
-            EntryPoints::APTPermissionedTransfer
-            | EntryPoints::APTTransfer => "permissioned_transfer",
+            EntryPoints::DeserializeU256 => "bcs_stream",
+            EntryPoints::APTPermissionedTransfer | EntryPoints::APTTransfer => {
+                "permissioned_transfer"
+            },
         }
     }
 
@@ -759,14 +760,12 @@ impl EntryPoints {
                     bcs::to_bytes(&1u64).unwrap(),
                 ],
             ),
-            EntryPoints::APTTransfer => get_payload(
-                module_id,
-                ident_str!("transfer").to_owned(),
-                vec![
+            EntryPoints::APTTransfer => {
+                get_payload(module_id, ident_str!("transfer").to_owned(), vec![
                     bcs::to_bytes(&other.expect("Must provide other")).unwrap(),
                     bcs::to_bytes(&1u64).unwrap(),
-                ],
-            )
+                ])
+            },
         }
     }
 
@@ -877,8 +876,9 @@ impl EntryPoints {
             EntryPoints::DeserializeU256 => AutomaticArgs::None,
             EntryPoints::IncGlobalMilestoneAggV2 { .. } => AutomaticArgs::None,
             EntryPoints::CreateGlobalMilestoneAggV2 { .. } => AutomaticArgs::Signer,
-            EntryPoints::APTPermissionedTransfer
-            | EntryPoints::APTTransfer => AutomaticArgs::Signer,
+            EntryPoints::APTPermissionedTransfer | EntryPoints::APTTransfer => {
+                AutomaticArgs::Signer
+            },
         }
     }
 }