chore: Refactor the op code gen to better support unions of input types

package.json

@@ -33,6 +33,7 @@
     "algo-ts:build:watch": "cd packages/algo-ts && npm run build:3-build:watch",
     "algo-ts:install": "npm i -D @algorandfoundation/algorand-typescript@./packages/algo-ts/dist",
     "npm-install-all": "cd packages/algo-ts && npm i && npm run build && cd ../../ && npm i",
+    "script-all": "run-p script:*",
     "script:op-types": "tsx ./scripts/generate-op-types.ts packages/algo-ts/src/op-types.ts && cd packages/algo-ts && eslint src/op-types.ts --fix",
     "script:op-metadata": "tsx ./scripts/generate-op-metadata.ts src/awst_build/op-metadata.ts && eslint src/awst_build/op-metadata.ts --fix",
     "script:op-ptypes": "tsx ./scripts/generate-op-ptypes.ts src/awst_build/ptypes/op-ptypes.ts && eslint src/awst_build/ptypes/op-ptypes.ts --fix"

packages/algo-ts/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@algorandfoundation/algorand-typescript",
-  "version": "0.0.1-alpha.18",
+  "version": "0.0.1-alpha.19",
   "description": "This package contains definitions for the types which comprise Algorand TypeScript which can be compiled to run on the Algorand Virtual Machine using the Puya compiler.",
   "private": false,
   "main": "index.js",

packages/algo-ts/src/op-types.ts

@@ -1,3 +1,7 @@
+/* THIS FILE IS GENERATED BY ~/scripts/generate-op-types.ts - DO NOT MODIFY DIRECTLY */
+import { bytes, uint64, biguint } from './primitives'
+import { Account, Application, Asset } from './reference'
+
 export enum Base64 {
   URLEncoding = 'URLEncoding',
   StdEncoding = 'StdEncoding',
@@ -3023,10 +3027,6 @@ export type VoterParamsType = {
 export type VrfVerifyType = (s: VrfVerify, a: bytes, b: bytes, c: bytes) => readonly [bytes, boolean]
 export type SelectType = ((a: bytes, b: bytes, c: uint64) => bytes) & ((a: uint64, b: uint64, c: uint64) => uint64)
 export type SetBitType = ((target: bytes, n: uint64, c: uint64) => bytes) & ((target: uint64, n: uint64, c: uint64) => uint64)
-/* THIS FILE IS GENERATED BY ~/scripts/generate-op-types.ts - DO NOT MODIFY DIRECTLY */
-import { biguint, bytes, uint64 } from './primitives'
-import { Account, Application, Asset } from './reference'
-
 export type OpsNamespace = {
   AcctParams: AcctParamsType
   addw: AddwType

scripts/build-op-module.ts

@@ -1,4 +1,5 @@
 import { camelCase, pascalCase } from 'change-case'
+import fs from 'fs'
 import langSpec from '../langspec.puya.json'
 import { hasFlags, invariant } from '../src/util'
 import type { Op } from './langspec'
@@ -114,6 +115,7 @@ const OPERATOR_OPCODES = new Set([
 ])
 
 export enum AlgoTsType {
+  None = 0,
   Bytes = 1 << 0,
   Uint64 = 1 << 1,
   Boolean = 1 << 2,
@@ -122,7 +124,8 @@ export enum AlgoTsType {
   Application = 1 << 5,
   Void = 1 << 6,
   BigUint = 1 << 7,
-  Enum = 1 << 8,
+  String = 1 << 8,
+  Enum = 1 << 9,
 }
 
 export type EnumValue = {
@@ -164,15 +167,12 @@ const TYPE_MAP: Record<string, AlgoTsType> = {
   any: AlgoTsType.Uint64 | AlgoTsType.Bytes,
 }
 
-const INPUT_TYPE_MAP: Record<string, AlgoTsType> = {
-  application: AlgoTsType.Application | AlgoTsType.Uint64,
-  asset: AlgoTsType.Asset | AlgoTsType.Uint64,
-}
-
 const INPUT_ALGOTS_TYPE_MAP: Record<AlgoTsType, AlgoTsType> = {
+  [AlgoTsType.None]: AlgoTsType.None,
   [AlgoTsType.Asset]: AlgoTsType.Asset | AlgoTsType.Uint64,
   [AlgoTsType.Application]: AlgoTsType.Application | AlgoTsType.Uint64,
   [AlgoTsType.Bytes]: AlgoTsType.Bytes,
+  [AlgoTsType.String]: AlgoTsType.String,
   [AlgoTsType.Uint64]: AlgoTsType.Uint64,
   [AlgoTsType.Boolean]: AlgoTsType.Boolean,
   [AlgoTsType.Account]: AlgoTsType.Account,
@@ -352,7 +352,7 @@ function isSplitableUnion(t: AlgoTsType): boolean {
   return !(hasFlags(t, AlgoTsType.Enum) || atomicTypes.includes(t))
 }
 /**
- * If a function returns a union type, split it into multiple functions for each part of the union
+ * If a function returns a union type, split into multiple functions for each part of the union
  *
  * eg.
  * get(): bytes | uint64
@@ -423,14 +423,14 @@ export function buildOpModule() {
           name: getEnumOpName(member.name, opNameConfig),
           immediateArgs: def.immediate_args.map((i) => ({
             name: camelCase(i.name),
-            type: getMappedType(i.immediate_type, i.arg_enum, true),
+            type: expandInputType(getMappedType(i.immediate_type, i.arg_enum)),
           })),
           stackArgs: def.stack_inputs.map((sa, i) => {
             if (i === enumArg.modifies_stack_input) {
               invariant(member.stackType, 'Member must have stack type')
-              return { name: camelCase(sa.name), type: INPUT_ALGOTS_TYPE_MAP[member.stackType] ?? member.stackType }
+              return { name: camelCase(sa.name), type: expandInputType(member.stackType) }
             }
-            return { name: camelCase(sa.name), type: getMappedType(sa.stack_type, null, true) }
+            return { name: camelCase(sa.name), type: expandInputType(getMappedType(sa.stack_type, null)) }
           }),
           returnTypes: def.stack_outputs.map((o, i) => {
             if (i === enumArg.modifies_stack_output) {
@@ -450,9 +450,9 @@ export function buildOpModule() {
           name: getOpName(def.name, opNameConfig),
           immediateArgs: def.immediate_args.map((i) => ({
             name: camelCase(i.name),
-            type: getMappedType(i.immediate_type, i.arg_enum, true),
+            type: expandInputType(getMappedType(i.immediate_type, i.arg_enum)),
           })),
-          stackArgs: def.stack_inputs.map((i) => ({ name: camelCase(i.name), type: getMappedType(i.stack_type, null, true) })),
+          stackArgs: def.stack_inputs.map((i) => ({ name: camelCase(i.name), type: expandInputType(getMappedType(i.stack_type, null)) })),
           returnTypes: def.stack_outputs.map((o) => getMappedType(o.stack_type, null)),
           docs: getOpDocs(def),
         }),
@@ -574,7 +574,7 @@ function getEnumOpName(enumMember: string, config: OpNameConfig): string {
   return camelCase([config.prefix, enumMember].filter(Boolean).join('_'))
 }
 
-function getMappedType(t: string | null, enumName: string | null, isInput: boolean = false): AlgoTsType {
+function getMappedType(t: string | null, enumName: string | null): AlgoTsType {
   invariant(t !== 'arg_enum' || enumName !== undefined, 'Must provide enumName for arg_enum types')
   if (t === null || t === undefined) {
     throw new Error('Missing type')
@@ -584,15 +584,34 @@ function getMappedType(t: string | null, enumName: string | null, isInput: boole
     invariant(enumDef, `Definition must exist for ${enumName}`)
     return enumDef.typeFlag
   }
-  const mappedType = isInput ? (INPUT_TYPE_MAP[t ?? ''] ?? TYPE_MAP[t ?? '']) : TYPE_MAP[t ?? '']
+  const mappedType = TYPE_MAP[t ?? '']
   invariant(mappedType, `Mapped type must exist for ${t}`)
   return mappedType
 }
 
+function splitBitFlags<T extends number>(aType: T): T[] {
+  if (!aType) return []
+  return new Array(Math.floor(Math.log2(aType)) + 1).fill(null).flatMap((_, i) => {
+    const n = (2 ** i) as T
+    return n & aType ? n : []
+  })
+}
+
+function expandInputType(aType: AlgoTsType): AlgoTsType {
+  return splitBitFlags(aType).reduce((acc, cur) => acc | (cur in INPUT_ALGOTS_TYPE_MAP ? INPUT_ALGOTS_TYPE_MAP[cur] : cur), AlgoTsType.None)
+}
+
 const getOpDocs = (op: Op): string[] => [
   ...(op.doc ?? [])
     .filter(Boolean)
     .map((d: string) => `${d.replace('params: ', '@param ').replace('Return: ', '\n * @return ')}`.split('\n').map((s) => s.trimEnd()))
     .flat(),
   `@see Native TEAL opcode: [\`${op.name}\`](https://developer.algorand.org/docs/get-details/dapps/avm/teal/opcodes/v10/#${op.name})`,
 ]
+
+function testOpModule() {
+  const mod = buildOpModule()
+
+  fs.writeFileSync('op-module.json', JSON.stringify(mod, undefined, 2))
+}
+testOpModule()

scripts/generate-op-metadata.ts

@@ -85,6 +85,10 @@ function* algoTsToPType(t: AlgoTsType) {
     t ^= AlgoTsType.Bytes
     yield 'ptypes.bytesPType'
   }
+  if (hasFlags(t, AlgoTsType.String)) {
+    t ^= AlgoTsType.String
+    yield 'ptypes.stringPType'
+  }
   if (Number(t) !== 0) throw new Error(`Unhandled flags ${t}`)
 }
 

scripts/generate-op-types.ts

@@ -7,7 +7,7 @@ import { AlgoTsType, buildOpModule, ENUMS_TO_EXPOSE } from './build-op-module'
 function* emitTypes(module: OpModule) {
   function* emitHeader() {
     yield `/* THIS FILE IS GENERATED BY ~/scripts/generate-op-types.ts - DO NOT MODIFY DIRECTLY */
-import { bytes, BytesCompat, uint64, Uint64Compat, biguint } from './primitives'
+import { bytes, uint64, biguint } from './primitives'
 import { Account, Application, Asset } from './reference'
 
 `
@@ -50,6 +50,7 @@ import { Account, Application, Asset } from './reference'
     if (hasFlags(returnType, AlgoTsType.Asset)) yield 'Asset'
     if (hasFlags(returnType, AlgoTsType.Uint64)) yield 'uint64'
     if (hasFlags(returnType, AlgoTsType.Bytes)) yield 'bytes'
+    if (hasFlags(returnType, AlgoTsType.String)) yield 'string'
     if (hasFlags(returnType, AlgoTsType.Boolean)) yield 'boolean'
     if (hasFlags(returnType, AlgoTsType.BigUint)) yield 'biguint'
     if (hasFlags(returnType, AlgoTsType.Void)) yield 'void'
@@ -83,15 +84,18 @@ import { Account, Application, Asset } from './reference'
     if (hasFlags(argType, AlgoTsType.Asset)) yield 'Asset'
     if (hasFlags(argType, AlgoTsType.Uint64)) yield 'uint64'
     if (hasFlags(argType, AlgoTsType.Bytes)) yield 'bytes'
+    if (hasFlags(argType, AlgoTsType.String)) yield 'string'
     if (hasFlags(argType, AlgoTsType.Boolean)) yield 'boolean'
     if (hasFlags(argType, AlgoTsType.BigUint)) yield 'biguint'
     if (hasFlags(argType, AlgoTsType.Void)) yield 'void'
     if (hasFlags(argType, AlgoTsType.Enum)) {
-      for (const enumDef of opModule.enums.filter((a) => hasFlags(a.typeFlag, argType))) {
+      for (const enumDef of opModule.enums.filter((a) => hasFlags(argType, a.typeFlag))) {
         yield enumDef.tsName
       }
     }
   }
+  yield* emitHeader()
+
   yield* emitEnums()
 
   for (const item of module.items) {
@@ -167,7 +171,6 @@ import { Account, Application, Asset } from './reference'
     }
   }
 
-  yield* emitHeader()
   yield `export type OpsNamespace = {\n`
   for (const item of opModule.items) {
     if (item.type === 'op-function' || item.type === 'op-overloaded-function') {