implement Scratch and add test for it

Author: boblat

src/impl/index.ts

@@ -8,3 +8,4 @@ export { GTxn } from './gtxn'
 export * from './pure'
 export { Txn, gaid } from './txn'
 export { GITxn, ITxn, ITxnCreate } from './itxn'
+export { Scratch } from './scratch'

src/impl/inner-transactions.ts

@@ -158,20 +158,20 @@ export class ApplicationInnerTxn extends ApplicationTransaction implements itxn.
   }
 }
 
-export const createInnerTxn = <TFields extends InnerTxnFields, T extends InnerTxn>(fields: TFields): T => {
+export const createInnerTxn = <TFields extends InnerTxnFields>(fields: TFields) => {
   switch (fields.type) {
     case TransactionType.Payment:
-      return new PaymentInnerTxn(fields) as T
+      return new PaymentInnerTxn(fields)
     case TransactionType.AssetConfig:
-      return new AssetConfigInnerTxn(fields) as T
+      return new AssetConfigInnerTxn(fields)
     case TransactionType.AssetTransfer:
-      return new AssetTransferInnerTxn(fields as itxn.AssetTransferFields) as T
+      return new AssetTransferInnerTxn(fields as itxn.AssetTransferFields)
     case TransactionType.AssetFreeze:
-      return new AssetFreezeInnerTxn(fields as itxn.AssetFreezeFields) as T
+      return new AssetFreezeInnerTxn(fields as itxn.AssetFreezeFields)
     case TransactionType.ApplicationCall:
-      return new ApplicationInnerTxn(fields) as unknown as T
+      return new ApplicationInnerTxn(fields)
     case TransactionType.KeyRegistration:
-      return new KeyRegistrationInnerTxn(fields) as T
+      return new KeyRegistrationInnerTxn(fields)
     default:
       throw new internal.errors.InternalError(`Invalid inner transaction type: ${fields.type}`)
   }
@@ -211,7 +211,7 @@ export class ItxnParams<TFields extends InnerTxnFields, TTransaction extends Inn
     this.#fields = { ...fields, type }
   }
   submit(): TTransaction {
-    return createInnerTxn<InnerTxnFields, TTransaction>(this.#fields)
+    return createInnerTxn<InnerTxnFields>(this.#fields) as unknown as TTransaction
   }
 
   set(p: Partial<TFields>) {

src/impl/scratch.ts

@@ -0,0 +1,25 @@
+import { bytes, internal, uint64 } from '@algorandfoundation/algo-ts'
+import { lazyContext } from '../context-helpers/internal-context'
+
+export const Scratch: internal.opTypes.ScratchType = {
+  loadBytes: function (a: internal.primitives.StubUint64Compat): bytes {
+    const result = lazyContext.activeGroup.activeTransaction.getScratchSlot(a)
+    if (result instanceof internal.primitives.BytesCls) {
+      return result as bytes
+    }
+    throw new internal.errors.InternalError('Invalid scratch slot type')
+  },
+  loadUint64: function (a: internal.primitives.StubUint64Compat): uint64 {
+    const result = lazyContext.activeGroup.activeTransaction.getScratchSlot(a)
+    if (result instanceof internal.primitives.Uint64Cls) {
+      return result as uint64
+    }
+    throw new internal.errors.InternalError('Invalid scratch slot type')
+  },
+  store: function (
+    a: internal.primitives.StubUint64Compat,
+    b: internal.primitives.StubUint64Compat | internal.primitives.StubBytesCompat,
+  ): void {
+    lazyContext.activeGroup.activeTransaction.setScratchSlot(a, b)
+  },
+}

src/impl/transactions.ts

@@ -14,7 +14,7 @@ import {
 import { MAX_ITEMS_IN_LOG } from '../constants'
 import { lazyContext } from '../context-helpers/internal-context'
 import { Mutable, ObjectKeys } from '../typescript-helpers'
-import { asBytes, asNumber, asUint64Cls, combineIntoMaxBytePages, getRandomBytes } from '../util'
+import { asBytes, asMaybeBytesCls, asMaybeUint64Cls, asNumber, asUint64Cls, combineIntoMaxBytePages, getRandomBytes } from '../util'
 
 const baseDefaultFields = () => ({
   sender: lazyContext.defaultSender,
@@ -44,6 +44,7 @@ abstract class TransactionBase {
     this.groupIndex = fields.groupIndex ?? baseDefaults.groupIndex
     this.txnId = fields.txnId ?? baseDefaults.txnId
     this.rekeyTo = fields.rekeyTo ?? baseDefaults.rekeyTo
+    this.scratchSpace = Array(256).fill(Uint64(0))
   }
 
   readonly sender: Account
@@ -56,6 +57,28 @@ abstract class TransactionBase {
   readonly groupIndex: uint64
   readonly txnId: bytes
   readonly rekeyTo: Account
+  readonly scratchSpace: Array<bytes | uint64>
+
+  setScratchSlot(
+    index: internal.primitives.StubUint64Compat,
+    value: internal.primitives.StubBytesCompat | internal.primitives.StubUint64Compat,
+  ): void {
+    const i = asNumber(index)
+    if (i >= this.scratchSpace.length) {
+      throw internal.errors.internalError('invalid scratch slot')
+    }
+    const bytesValue = asMaybeBytesCls(value)
+    const uint64Value = asMaybeUint64Cls(value)
+    this.scratchSpace[i] = bytesValue?.asAlgoTs() ?? uint64Value?.asAlgoTs() ?? Uint64(0)
+  }
+
+  getScratchSlot(index: internal.primitives.StubUint64Compat): bytes | uint64 {
+    const i = asNumber(index)
+    if (i >= this.scratchSpace.length) {
+      throw internal.errors.internalError('invalid scratch slot')
+    }
+    return this.scratchSpace[i]
+  }
 }
 
 export class PaymentTransaction extends TransactionBase implements gtxn.PaymentTxn {
@@ -201,6 +224,7 @@ export type ApplicationTransactionFields = TxnFields<gtxn.ApplicationTxn> &
     approvalProgramPages: Array<bytes>
     clearStateProgramPages: Array<bytes>
     appLogs: Array<bytes>
+    scratchSpace: Array<bytes | uint64>
   }>
 
 export class ApplicationTransaction extends TransactionBase implements gtxn.ApplicationTxn {
@@ -233,6 +257,7 @@ export class ApplicationTransaction extends TransactionBase implements gtxn.Appl
     this.#apps = fields.apps ?? []
     this.#approvalProgramPages = fields.approvalProgramPages ?? (fields.approvalProgram ? [fields.approvalProgram] : [])
     this.#clearStateProgramPages = fields.clearStateProgramPages ?? (fields.clearStateProgram ? [fields.clearStateProgram] : [])
+    fields.scratchSpace?.forEach((v, i) => this.setScratchSlot(i, v))
   }
 
   readonly appId: Application

src/subcontexts/transaction-context.ts

@@ -1,4 +1,4 @@
-import { internal, TransactionType, uint64 } from '@algorandfoundation/algo-ts'
+import { bytes, internal, TransactionType, uint64 } from '@algorandfoundation/algo-ts'
 import algosdk from 'algosdk'
 import { lazyContext } from '../context-helpers/internal-context'
 import { DecodedLogs, decodeLogs, LogDecoding } from '../decode-logs'
@@ -162,6 +162,10 @@ export class TransactionGroup {
     return this.constructingItxnGroup.at(-1)!
   }
 
+  getScratchSlot(index: internal.primitives.StubUint64Compat): bytes | uint64 {
+    return this.activeTransaction.getScratchSlot(index)
+  }
+
   patchActiveTransactionFields(fields: AllTransactionFields) {
     const activeTransaction = this.activeTransaction as unknown as AllTransactionFields
     const filteredFields = Object.fromEntries(Object.entries(fields).filter(([_, value]) => value !== undefined))

tests/state-op-codes.spec.ts

@@ -367,4 +367,23 @@ describe('State op codes', async () => {
       })
     })
   })
+
+  describe('Scratch', async () => {
+    test.each([
+      [0, Bytes('test_bytes')],
+      [1, Uint64(42)],
+      [2, Bytes('test_bytes')],
+      [3, Uint64(42)],
+      [255, Bytes('max_index')],
+    ])('should return the correct field value of the scratch slot', async (index: number, value: bytes | uint64) => {
+      const newScratchSpace = Array(256).fill(Uint64(0))
+      newScratchSpace[index] = value
+
+      ctx.txn.createScope([ctx.any.txn.applicationCall({ scratchSpace: newScratchSpace })]).execute(() => {})
+
+      expect(ctx.txn.lastGroup.getScratchSlot(index)).toEqual(value)
+
+      expect(() => ctx.txn.lastGroup.getScratchSlot(256)).toThrow('invalid scratch slot')
+    })
+  })
 })