fix(sui): view function address should honor provided address (#96)

sentioxyz · May 11, 2024 · 69d4fcf · 69d4fcf
1 parent 5828381
commit 69d4fcf
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 21 deletions.
diff --git a/packages/aptos/src/codegen/codegen.ts b/packages/aptos/src/codegen/codegen.ts
@@ -47,7 +47,7 @@ export class AptosCodegen extends AbstractCodegen<MoveModuleBytecode, Event | Mo
       import { Aptos, Account as AptosAccount, MoveAddressType, PendingTransactionResponse, InputGenerateTransactionOptions, MoveStructId, InputViewFunctionData } from '@aptos-labs/ts-sdk'
     `
   }
-  protected generateExtra(module: InternalMoveModule) {
+  protected generateExtra(address: string, module: InternalMoveModule) {
     const funcs = module.exposedFunctions.map((f) => this.generateEntryForFunction(module, f))
 
     const viewFuncs = module.exposedFunctions.map((f) => this.generateViewFunction(module, f))

diff --git a/packages/move/src/abstract-codegen.ts b/packages/move/src/abstract-codegen.ts
@@ -85,11 +85,23 @@ export abstract class AbstractCodegen<ModuleTypes, StructType> {
       for (const module of modules) {
         loader.register(module, path.basename(file, '.json'))
       }
-      const codeGen = new AccountCodegen(this, loader, abi, modules, {
-        fileName: path.basename(file, '.json'),
-        outputDir: outputDir
-        // network,
-      })
+
+      const fileName = path.basename(file, '.json')
+      // TODO more reliable way to check if this is sui address
+      const address = fileName.startsWith('0x') ? fileName : undefined
+
+      const codeGen = new AccountCodegen(
+        this,
+        loader,
+        abi,
+        modules,
+        {
+          fileName,
+          outputDir: outputDir
+          // network,
+        },
+        address
+      )
 
       outputs.push(...codeGen.generate())
     }
@@ -110,11 +122,18 @@ export abstract class AbstractCodegen<ModuleTypes, StructType> {
           for (const module of modules) {
             loader.register(module, account)
           }
-          const codeGen = new AccountCodegen(this, loader, rawModules, modules, {
-            fileName: account,
-            outputDir: outputDir
-            // network,
-          })
+          const codeGen = new AccountCodegen(
+            this,
+            loader,
+            rawModules,
+            modules,
+            {
+              fileName: account,
+              outputDir: outputDir
+              // network,
+            },
+            account
+          )
 
           outputs.push(...codeGen.generate())
         } catch (e) {
@@ -149,13 +168,14 @@ export abstract class AbstractCodegen<ModuleTypes, StructType> {
     return outputs.length
   }
 
-  protected generateExtra(module: InternalMoveModule) {
+  protected generateExtra(address: string | undefined, module: InternalMoveModule) {
     return ''
   }
 
   generateModule(
     module: InternalMoveModule,
-    allEventStructs: Map<string, InternalMoveStruct>
+    allEventStructs: Map<string, InternalMoveStruct>,
+    addressOverride?: string
     // network: NetworkType
   ) {
     const qname = moduleQname(module)
@@ -179,7 +199,7 @@ export abstract class AbstractCodegen<ModuleTypes, StructType> {
   export namespace ${moduleName} {
     ${structs.join('\n')}
     
-    ${this.generateExtra(module)}
+    ${this.generateExtra(addressOverride, module)}
     
     ${events.join('\n')}
     
@@ -409,20 +429,24 @@ export class AccountCodegen<ModuleType, StructType> {
   abi: ModuleType[]
   loader: AccountRegister
   moduleGen: AbstractCodegen<ModuleType, StructType>
+  // Usually it's same as module.address, but in upgraded package, this might be the new address
+  address?: string
 
   constructor(
     moduleGen: AbstractCodegen<ModuleType, StructType>,
     loader: AccountRegister,
     abi: ModuleType[],
     modules: InternalMoveModule[],
-    config: Config
+    config: Config,
+    address: string | undefined
   ) {
     // const json = fs.readFileSync(config.srcFile, 'utf-8')
     this.moduleGen = moduleGen
     this.abi = abi
     this.modules = modules
     this.config = config
     this.loader = loader
+    this.address = address
   }
 
   generate(): OutputFile[] {
@@ -470,13 +494,13 @@ export class AccountCodegen<ModuleType, StructType> {
     /* tslint:disable */
     /* eslint-disable */
 
-    /* Generated modules for account ${address} */
+    /* Generated types for ${this.address || address}, original address ${address} */
 
     ${this.moduleGen.generateImports()}
 
     ${moduleImports.join('\n')}
 
-    ${this.modules.map((m) => this.moduleGen.generateModule(m, eventsMap)).join('\n')}
+    ${this.modules.map((m) => this.moduleGen.generateModule(m, eventsMap, this.address)).join('\n')}
 
     const MODULES = JSON.parse('${JSON.stringify(this.abi)}')
 

diff --git a/packages/sui/src/codegen/codegen.ts b/packages/sui/src/codegen/codegen.ts
@@ -104,8 +104,10 @@ export class SuiCodegen extends AbstractCodegen<
     return super.generateForEvents(module, struct)
   }
 
-  protected generateExtra(module: InternalMoveModule): string {
-    const funcs = module.exposedFunctions.map((f) => this.generateBuilderForFunction(module, f))
+  protected generateExtra(address: string | undefined, module: InternalMoveModule): string {
+    const funcs = module.exposedFunctions.map((f) =>
+      this.generateBuilderForFunction(address || module.address, module, f)
+    )
 
     const viewFuncs = module.exposedFunctions.map((f) => this.generateViewFunction(module, f))
 
@@ -177,7 +179,11 @@ export class SuiCodegen extends AbstractCodegen<
     }`
   }
 
-  protected generateBuilderForFunction(module: InternalMoveModule, func: InternalMoveFunction): string {
+  protected generateBuilderForFunction(
+    address: string,
+    module: InternalMoveModule,
+    func: InternalMoveFunction
+  ): string {
     if (func.visibility === InternalMoveFunctionVisibility.PRIVATE && func.isEntry !== true) {
       return ''
     }
@@ -206,7 +212,7 @@ export class SuiCodegen extends AbstractCodegen<
       
       // @ts-ignore
       return tx.moveCall({
-        target: "${module.address}::${module.name}::${func.name}",
+        target: "${address}::${module.name}::${func.name}",
         arguments: _args,
         ${typeParamArg.length > 0 ? `typeArguments: [${typeParamToString}]` : ``}
       })