feat: implement stubs for appOptedIn, ensureBudget, TemplateVar

refactor: function call replacement (...Impl) in visitor to be more
resilient by catering for the following cases
  - module name prefix
    (e.g. `arc4.interpretAsArc4`)
  - import name alias
    (e.g. `import { interpretAsArc4 as interpret } from '@algorandfoundation/algorand-typescript/arc4'`)
  - function name conflict
    (e.g.
    ```
    import { interpretAsArc4 as interpret } ....
    const interpretAsArc4 = () => {...}
    ```)
feat: add arc4 value generator
feat: add zk-whitelist example contract and tests

Author: boblat

examples/rollup.config.ts

@@ -12,6 +12,7 @@ const config: RollupOptions = {
     'examples/auction/contract.algo.ts',
     'examples/voting/contract.algo.ts',
     'examples/simple-voting/contract.algo.ts',
+    'examples/zk-whitelist/contract.algo.ts',
   ],
   output: [
     {

examples/zk-whitelist/contract.algo.ts

@@ -0,0 +1,94 @@
+import {
+  abimethod,
+  Account,
+  arc4,
+  assert,
+  BigUint,
+  Bytes,
+  ensureBudget,
+  Global,
+  itxn,
+  LocalState,
+  op,
+  OpUpFeeSource,
+  TemplateVar,
+  Txn,
+  uint64,
+} from '@algorandfoundation/algorand-typescript'
+
+const curveMod = 21888242871839275222246405745257275088548364400416034343698204186575808495617n
+const verifierBudget = 145000
+
+export default class ZkWhitelistContract extends arc4.Contract {
+  appName: arc4.Str | undefined
+  whiteList = LocalState<boolean>()
+
+  @abimethod({ onCreate: 'require' })
+  create(name: arc4.Str) {
+    // Create the application
+    this.appName = name
+  }
+
+  @abimethod({ allowActions: ['UpdateApplication', 'DeleteApplication'] })
+  update() {
+    // Update the application if it is mutable (manager only)
+    assert(Global.creatorAddress === Txn.sender)
+  }
+
+  @abimethod({ allowActions: ['OptIn', 'CloseOut'] })
+  optInOrOut() {
+    // Opt in or out of the application
+    return
+  }
+
+  @abimethod()
+  addAddressToWhitelist(address: arc4.Address, proof: arc4.DynamicArray<arc4.Address>): arc4.Str {
+    /*
+    Add caller to the whitelist if the zk proof is valid.
+    On success, will return an empty string. Otherwise, will return an error
+    message.
+    */
+    ensureBudget(verifierBudget, OpUpFeeSource.GroupCredit)
+    // The verifier expects public inputs to be in the curve field, but an
+    // Algorand address might represent a number larger than the field
+    // modulus, so to be safe we take the address modulo the field modulus
+    const addressMod = arc4.interpretAsArc4<arc4.Address>(op.bzero(32).bitwiseOr(Bytes(BigUint(address.bytes) % curveMod)))
+    // Verify the proof by calling the deposit verifier app
+    const verified = this.verifyProof(TemplateVar<uint64>('VERIFIER_APP_ID'), proof, new arc4.DynamicArray(addressMod))
+    if (!verified.native) {
+      return new arc4.Str('Proof verification failed')
+    }
+    // if successful, add the sender to the whitelist by setting local state
+    const account = Account(address.bytes)
+    if (Txn.sender !== account) {
+      return new arc4.Str('Sender address does not match authorized address')
+    }
+    this.whiteList(account).value = true
+    return new arc4.Str('')
+  }
+
+  @abimethod()
+  isOnWhitelist(address: arc4.Address): arc4.Bool {
+    // Check if an address is on the whitelist
+    const account = address.native
+    const optedIn = op.appOptedIn(account, Global.currentApplicationId)
+    if (!optedIn) {
+      return new arc4.Bool(false)
+    }
+    const whitelisted = this.whiteList(account).value
+    return new arc4.Bool(whitelisted)
+  }
+
+  verifyProof(appId: uint64, proof: arc4.DynamicArray<arc4.Address>, publicInputs: arc4.DynamicArray<arc4.Address>): arc4.Bool {
+    // Verify a proof using the verifier app.
+    const verified = itxn
+      .applicationCall({
+        appId: appId,
+        fee: 0,
+        appArgs: [arc4.methodSelector('verify(byte[32][],byte[32][])bool'), proof.copy(), publicInputs.copy()],
+        onCompletion: arc4.OnCompleteAction.NoOp,
+      })
+      .submit().lastLog
+    return arc4.interpretAsArc4<arc4.Bool>(verified, 'log')
+  }
+}

examples/zk-whitelist/contract.spec.ts

@@ -0,0 +1,79 @@
+import { arc4, Bytes } from '@algorandfoundation/algorand-typescript'
+import { TestExecutionContext } from '@algorandfoundation/algorand-typescript-testing'
+import { afterEach, describe, expect, it } from 'vitest'
+import { ABI_RETURN_VALUE_LOG_PREFIX } from '../../src/constants'
+import ZkWhitelistContract from './contract.algo'
+
+describe('ZK Whitelist', () => {
+  const ctx = new TestExecutionContext()
+
+  afterEach(() => {
+    ctx.reset()
+  })
+
+  it('should be able to add address to whitelist', () => {
+    // Arrange
+    const contract = ctx.contract.create(ZkWhitelistContract)
+    contract.create(ctx.any.arc4.str(10))
+
+    const address = new arc4.Address(ctx.defaultSender)
+    const proof = new arc4.DynamicArray<arc4.Address>(new arc4.Address(Bytes(new Uint8Array(Array(32).fill(0)))))
+
+    const dummyVerifierApp = ctx.any.application({ appLogs: [ABI_RETURN_VALUE_LOG_PREFIX.concat(Bytes.fromHex('80'))] })
+    ctx.setTemplateVar('VERIFIER_APP_ID', dummyVerifierApp.id)
+
+    // Act
+    const result = contract.addAddressToWhitelist(address, proof)
+
+    // Assert
+    expect(result.native).toEqual('')
+    expect(contract.whiteList(ctx.defaultSender).value).toEqual(true)
+  })
+
+  it('returns error message if proof verification fails', () => {
+    // Arrange
+    const contract = ctx.contract.create(ZkWhitelistContract)
+    contract.create(ctx.any.arc4.str(10))
+
+    const address = ctx.any.arc4.address()
+    const proof = new arc4.DynamicArray<arc4.Address>(new arc4.Address(Bytes(new Uint8Array(Array(32).fill(0)))))
+    const dummyVerifierApp = ctx.any.application({ appLogs: [ABI_RETURN_VALUE_LOG_PREFIX.concat(Bytes(''))] })
+    ctx.setTemplateVar('VERIFIER_APP_ID', dummyVerifierApp.id)
+
+    // Act
+    const result = contract.addAddressToWhitelist(address, proof)
+
+    // Assert
+    expect(result.native).toEqual('Proof verification failed')
+  })
+
+  it('returns true if address is already in whitelist', () => {
+    // Arrange
+    const contract = ctx.contract.create(ZkWhitelistContract)
+    contract.create(ctx.any.arc4.str(10))
+
+    const dummyAccount = ctx.any.account({ optedApplications: [ctx.ledger.getApplicationForContract(contract)] })
+    contract.whiteList(dummyAccount).value = true
+
+    // Act
+    const result = contract.isOnWhitelist(new arc4.Address(dummyAccount))
+
+    // Assert
+    expect(result.native).toBe(true)
+  })
+
+  it('returns false if address is not in whitelist', () => {
+    // Arrange
+    const contract = ctx.contract.create(ZkWhitelistContract)
+    contract.create(ctx.any.arc4.str(10))
+
+    const dummyAccount = ctx.any.account({ optedApplications: [ctx.ledger.getApplicationForContract(contract)] })
+    contract.whiteList(dummyAccount).value = false
+
+    // Act
+    const result = contract.isOnWhitelist(new arc4.Address(dummyAccount))
+
+    // Assert
+    expect(result.native).toBe(false)
+  })
+})

src/constants.ts

@@ -3,7 +3,11 @@ import { Bytes } from '@algorandfoundation/algorand-typescript'
 export const UINT64_SIZE = 64
 export const UINT512_SIZE = 512
 export const MAX_UINT8 = 2 ** 8 - 1
+export const MAX_UINT16 = 2 ** 16 - 1
+export const MAX_UINT32 = 2 ** 32 - 1
 export const MAX_UINT64 = 2n ** 64n - 1n
+export const MAX_UINT128 = 2n ** 128n - 1n
+export const MAX_UINT256 = 2n ** 256n - 1n
 export const MAX_UINT512 = 2n ** 512n - 1n
 export const MAX_BYTES_SIZE = 4096
 export const MAX_LOG_SIZE = 1024
@@ -43,6 +47,7 @@ export const ABI_RETURN_VALUE_LOG_PREFIX = Bytes.fromHex('151F7C75')
 
 export const UINT64_OVERFLOW_UNDERFLOW_MESSAGE = 'Uint64 overflow or underflow'
 export const BIGUINT_OVERFLOW_UNDERFLOW_MESSAGE = 'BigUint overflow or underflow'
+export const DEFAULT_TEMPLATE_VAR_PREFIX = 'TMPL_'
 
 export const APP_ID_PREFIX = 'appID'
 export const HASH_BYTES_LENGTH = 32

src/impl/acct-params.ts

@@ -1,6 +1,7 @@
-import { Account, gtxn, internal, uint64 } from '@algorandfoundation/algorand-typescript'
+import { Account, Application, gtxn, internal, uint64 } from '@algorandfoundation/algorand-typescript'
 import { lazyContext } from '../context-helpers/internal-context'
 import { asMaybeUint64Cls } from '../util'
+import { getApp } from './app-params'
 
 export const getAccount = (acct: Account | internal.primitives.StubUint64Compat): Account => {
   const acctId = asMaybeUint64Cls(acct)
@@ -21,6 +22,19 @@ export const minBalance = (a: Account | internal.primitives.StubUint64Compat): u
   return acct.minBalance
 }
 
+export const appOptedIn = (
+  a: Account | internal.primitives.StubUint64Compat,
+  b: Application | internal.primitives.StubUint64Compat,
+): boolean => {
+  const account = getAccount(a)
+  const app = getApp(b)
+
+  if (account === undefined || app === undefined) {
+    return false
+  }
+  return account.isOptedIn(app)
+}
+
 export const AcctParams: internal.opTypes.AcctParamsType = {
   acctBalance(a: Account | internal.primitives.StubUint64Compat): readonly [uint64, boolean] {
     const acct = getAccount(a)

src/impl/ensure-budget.ts

@@ -0,0 +1,5 @@
+import { OpUpFeeSource, uint64 } from '@algorandfoundation/algorand-typescript'
+
+export function ensureBudgetImpl(_budget: uint64, _feeSource: OpUpFeeSource = OpUpFeeSource.GroupCredit) {
+  // ensureBudget function is emulated to be a no-op
+}

src/impl/index.ts

@@ -1,4 +1,4 @@
-export { AcctParams, balance, minBalance } from './acct-params'
+export { AcctParams, appOptedIn, balance, minBalance } from './acct-params'
 export { AppGlobal } from './app-global'
 export { AppLocal } from './app-local'
 export { AppParams } from './app-params'

src/impl/template-var.ts

@@ -0,0 +1,11 @@
+import { internal } from '@algorandfoundation/algorand-typescript'
+import { DEFAULT_TEMPLATE_VAR_PREFIX } from '../constants'
+import { lazyContext } from '../context-helpers/internal-context'
+
+export function TemplateVarImpl<T>(variableName: string, prefix = DEFAULT_TEMPLATE_VAR_PREFIX): T {
+  const key = prefix + variableName
+  if (!Object.hasOwn(lazyContext.value.templateVars, key)) {
+    throw internal.errors.codeError(`Template variable ${key} not found in test context!`)
+  }
+  return lazyContext.value.templateVars[prefix + variableName] as T
+}

src/runtime-helpers.ts

@@ -8,6 +8,8 @@ import { nameOfType } from './util'
 
 export { attachAbiMetadata } from './abi-metadata'
 export * from './impl/encoded-types'
+export { ensureBudgetImpl } from './impl/ensure-budget'
+export { TemplateVarImpl } from './impl/template-var'
 
 export function switchableValue(x: unknown): bigint | string | boolean {
   if (typeof x === 'boolean') return x
@@ -92,6 +94,7 @@ function arc4EncodedOp(left: ARC4Encoded, right: ARC4Encoded, op: BinaryOps): De
 function accountBinaryOp(left: AccountCls, right: AccountCls, op: BinaryOps): DeliberateAny {
   switch (op) {
     case '===':
+    case '!==':
       return bytesBinaryOp(left.bytes, right.bytes, op)
     default:
       internal.errors.internalError(`Unsupported operator ${op}`)

src/test-execution-context.ts

@@ -1,5 +1,6 @@
-import { Account, Application, Asset, Bytes, bytes, internal, uint64 } from '@algorandfoundation/algorand-typescript'
+import { Account, Application, Asset, bytes, internal, LogicSig, uint64 } from '@algorandfoundation/algorand-typescript'
 import { captureMethodConfig } from './abi-metadata'
+import { DEFAULT_TEMPLATE_VAR_PREFIX } from './constants'
 import { DecodedLogs, LogDecoding } from './decode-logs'
 import * as ops from './impl'
 import { AccountCls } from './impl/account'
@@ -18,6 +19,7 @@ import { Box, BoxMap, BoxRef, GlobalState, LocalState } from './impl/state'
 import { ContractContext } from './subcontexts/contract-context'
 import { LedgerContext } from './subcontexts/ledger-context'
 import { TransactionContext } from './subcontexts/transaction-context'
+import { DeliberateAny } from './typescript-helpers'
 import { getRandomBytes } from './util'
 import { ValueGenerator } from './value-generators'
 
@@ -28,6 +30,7 @@ export class TestExecutionContext implements internal.ExecutionContext {
   #valueGenerator: ValueGenerator
   #defaultSender: Account
   #activeLogicSigArgs: bytes[]
+  #template_vars: Record<string, DeliberateAny> = {}
 
   constructor(defaultSenderAddress?: bytes) {
     internal.ctxMgr.instance = this
@@ -126,6 +129,10 @@ export class TestExecutionContext implements internal.ExecutionContext {
     return this.#activeLogicSigArgs
   }
 
+  get templateVars(): Record<string, DeliberateAny> {
+    return this.#template_vars
+  }
+
   executeLogicSig(logicSig: LogicSig, ...args: bytes[]): boolean | uint64 {
     this.#activeLogicSigArgs = args
     try {
@@ -135,10 +142,16 @@ export class TestExecutionContext implements internal.ExecutionContext {
     }
   }
 
+  setTemplateVar(name: string, value: DeliberateAny) {
+    this.#template_vars[DEFAULT_TEMPLATE_VAR_PREFIX + name] = value
+  }
+
   reset() {
     this.#contractContext = new ContractContext()
     this.#ledgerContext = new LedgerContext()
     this.#txnContext = new TransactionContext()
+    this.#activeLogicSigArgs = []
+    this.#template_vars = {}
     internal.ctxMgr.reset()
     internal.ctxMgr.instance = this
   }