Merge pull request #25 from algorandfoundation/refactor-visitor

refactor: capture generic type info only for boxes, global and local states
algorandfoundation · Jan 3, 2025 · dafba04 · dafba04
2 parents 7b195f7 + 6d05031
commit dafba04
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 15 deletions.
diff --git a/src/test-execution-context.ts b/src/test-execution-context.ts
@@ -64,7 +64,6 @@ export class TestExecutionContext implements internal.ExecutionContext {
     this.txn.appendLog(value)
   }
 
-  /* @internal */
   exportLogs<const T extends [...LogDecoding[]]>(appId: uint64, ...decoding: T): DecodedLogs<T> {
     return this.txn.exportLogs(appId, ...decoding)
   }

diff --git a/src/test-transformer/node-factory.ts b/src/test-transformer/node-factory.ts
@@ -94,19 +94,17 @@ export const nodeFactory = {
     )
   },
 
-  callStubbedFunction(functionName: string, node: ts.CallExpression, ...typeInfos: (TypeInfo | undefined)[]) {
-    const infoStringArray = typeInfos.length ? typeInfos.map((typeInfo) => JSON.stringify(typeInfo)) : undefined
+  callStubbedFunction(functionName: string, node: ts.CallExpression, typeInfo?: TypeInfo) {
+    const typeInfoArg = typeInfo ? factory.createStringLiteral(JSON.stringify(typeInfo)) : undefined
     const updatedPropertyAccessExpression = factory.createPropertyAccessExpression(
       factory.createIdentifier('runtimeHelpers'),
       `${functionName}Impl`,
     )
-    const typeInfoArgs = infoStringArray
-      ? infoStringArray?.filter((s) => !!s).map((infoString) => factory.createStringLiteral(infoString))
-      : undefined
+
     return factory.createCallExpression(
       updatedPropertyAccessExpression,
       node.typeArguments,
-      [...(typeInfoArgs ?? []), ...(node.arguments ?? [])].filter((arg) => !!arg),
+      [typeInfoArg, ...(node.arguments ?? [])].filter((arg) => !!arg),
     )
   },
 } satisfies Record<string, (...args: DeliberateAny[]) => ts.Node>
diff --git a/src/test-transformer/visitors.ts b/src/test-transformer/visitors.ts
@@ -90,6 +90,7 @@ class ExpressionVisitor {
     private context: ts.TransformationContext,
     private helper: VisitorHelper,
     private expressionNode: ts.Expression,
+    private stubbedFunctionName?: string,
   ) {}
 
   public result(): ts.Expression {
@@ -104,6 +105,7 @@ class ExpressionVisitor {
       if (type instanceof ptypes.FunctionPType) type = type.returnType
 
       const isGeneric = isGenericType(type)
+      const needsToCaptureTypeInfo = isGeneric && isStateOrBoxType(type)
       const isArc4Encoded = isArc4EncodedType(type)
       const info = isGeneric || isArc4Encoded ? getGenericTypeInfo(type) : undefined
       let updatedNode = node
@@ -113,20 +115,22 @@ class ExpressionVisitor {
           updatedNode = nodeFactory.instantiateARC4EncodedType(updatedNode, info)
         }
       }
+
       if (ts.isCallExpression(updatedNode)) {
-        const stubbedFunctionName = tryGetStubbedFunctionName(updatedNode, this.helper)
-        const infos = [info]
+        const stubbedFunctionName = this.stubbedFunctionName ?? tryGetStubbedFunctionName(updatedNode, this.helper)
+        this.stubbedFunctionName = undefined
+        let infoArg = info
         if (isCallingEmit(stubbedFunctionName)) {
-          infos[0] = this.helper.resolveTypeParameters(updatedNode).map(getGenericTypeInfo)[0]
+          infoArg = this.helper.resolveTypeParameters(updatedNode).map(getGenericTypeInfo)[0]
         }
         if (isCallingDecodeArc4(stubbedFunctionName)) {
           const targetType = ptypes.ptypeToArc4EncodedType(type, this.helper.sourceLocation(node))
           const targetTypeInfo = getGenericTypeInfo(targetType)
-          infos[0] = targetTypeInfo
+          infoArg = targetTypeInfo
         }
-        updatedNode = stubbedFunctionName ? nodeFactory.callStubbedFunction(stubbedFunctionName, updatedNode, ...infos) : updatedNode
+        updatedNode = stubbedFunctionName ? nodeFactory.callStubbedFunction(stubbedFunctionName, updatedNode, infoArg) : updatedNode
       }
-      return isGeneric
+      return needsToCaptureTypeInfo
         ? nodeFactory.captureGenericTypeInfo(ts.visitEachChild(updatedNode, this.visit, this.context), JSON.stringify(info))
         : ts.visitEachChild(updatedNode, this.visit, this.context)
     }
@@ -218,8 +222,11 @@ class FunctionOrMethodVisitor {
     if (ts.isNewExpression(node)) {
       return new ExpressionVisitor(this.context, this.helper, node).result()
     }
-    if (ts.isCallExpression(node) && tryGetStubbedFunctionName(node, this.helper)) {
-      return new ExpressionVisitor(this.context, this.helper, node).result()
+    if (ts.isCallExpression(node)) {
+      const stubbedFunctionName = tryGetStubbedFunctionName(node, this.helper)
+      if (stubbedFunctionName) {
+        return new ExpressionVisitor(this.context, this.helper, node, stubbedFunctionName).result()
+      }
     }
 
     return node
@@ -305,6 +312,9 @@ const isGenericType = (type: ptypes.PType): boolean =>
     ptypes.TuplePType,
   )
 
+const isStateOrBoxType = (type: ptypes.PType): boolean =>
+  instanceOfAny(type, ptypes.BoxMapPType, ptypes.BoxPType, ptypes.GlobalStateType, ptypes.LocalStateType)
+
 const isArc4EncodedType = (type: ptypes.PType): boolean =>
   instanceOfAny(
     type,