feat: Reserved scratch space

package.json

@@ -9,7 +9,7 @@
     "dev:examples": "tsx src/cli.ts build examples --output-awst --output-awst-json",
     "dev:approvals": "rimraf tests/approvals/out && tsx src/cli.ts build tests/approvals --dry-run",
     "dev:expected-output": "tsx src/cli.ts build tests/expected-output --dry-run",
-    "dev:testing": "tsx src/cli.ts build tests/approvals/state-totals.algo.ts --output-awst --output-awst-json --output-ssa-ir --log-level=info --out-dir out/[name] --optimization-level=0",
+    "dev:testing": "tsx src/cli.ts build tests/approvals/reserve-scratch.algo.ts --output-awst --output-awst-json --output-ssa-ir --log-level=info --out-dir out/[name] --optimization-level=0",
     "audit": "better-npm-audit audit",
     "format": "prettier --write .",
     "lint": "eslint \"src/**/*.ts\"",

src/awst/to-code-visitor.ts

@@ -1,6 +1,5 @@
 import { Buffer } from 'node:buffer'
 import { TodoError } from '../errors'
-import { logger } from '../logger'
 import { uint8ArrayToBase32, uint8ArrayToUtf8 } from '../util'
 import type { ContractReference } from './models'
 import type {
@@ -408,7 +407,7 @@ export class ToCodeVisitor
       }
     }
     if (c.reservedScratchSpace.size) {
-      logger.warn(c.sourceLocation, 'Handle reservedScratchSpace to-code')
+      body.push(`reservedScratchSpace: [${Array.from(c.reservedScratchSpace).join(', ')}]`)
     }
     if (c.approvalProgram) {
       body.push(...this.visitSpecialMethod(c.approvalProgram, 'approvalProgram'))

src/awst_build/eb/contract-builder.ts

@@ -24,7 +24,9 @@ import { instanceEb } from '../type-registry'
 import { BaseContractMethodExpressionBuilder, ContractMethodExpressionBuilder } from './free-subroutine-expression-builder'
 import type { NodeBuilder } from './index'
 import { DecoratorDataBuilder, FunctionBuilder, InstanceBuilder } from './index'
-import { requireIntegerConstant, requireStringConstant } from './util'
+import { ArrayLiteralExpressionBuilder } from './literal/array-literal-expression-builder'
+import { BigIntLiteralExpressionBuilder } from './literal/big-int-literal-expression-builder'
+import { requireStringConstant } from './util'
 import { parseFunctionArgs } from './util/arg-parsing'
 import { requireAvmVersion } from './util/avm-version'
 import { VoidExpressionBuilder } from './void-expression-builder'
@@ -127,7 +129,7 @@ export class ContractOptionsDecoratorBuilder extends FunctionBuilder {
   readonly ptype = contractOptionsDecorator
   call(args: ReadonlyArray<NodeBuilder>, typeArgs: ReadonlyArray<PType>, sourceLocation: SourceLocation): NodeBuilder {
     const {
-      args: [{ avmVersion, name, stateTotals }],
+      args: [{ avmVersion, name, stateTotals, scratchSlots }],
     } = parseFunctionArgs({
       args,
       typeArgs,
@@ -149,15 +151,52 @@ export class ContractOptionsDecoratorBuilder extends FunctionBuilder {
       avmVersion: avmVersion && requireAvmVersion(avmVersion),
       name: name && requireStringConstant(name).value,
       stateTotals: stateTotals && buildStateTotals(stateTotals),
+      scratchSlots: scratchSlots && processScratchRanges(scratchSlots),
       sourceLocation,
     })
   }
 }
 
+function getRangeProp(builder: NodeBuilder, name: string): bigint {
+  if (builder.hasProperty(name)) {
+    return getLiteralNumber(builder.memberAccess(name, builder.sourceLocation))
+  }
+  throw new CodeError('Scratch slot reservations should be either a single slot or an object containing a from and to property', {
+    sourceLocation: builder.sourceLocation,
+  })
+}
+
+function getLiteralNumber(builder: NodeBuilder) {
+  codeInvariant(builder instanceof BigIntLiteralExpressionBuilder, 'Expected numeric literal', builder.sourceLocation)
+  return builder.value
+}
+
+function processScratchRanges(builder: NodeBuilder): Set<bigint> {
+  codeInvariant(
+    builder instanceof ArrayLiteralExpressionBuilder,
+    'Scratch ranges should be specified in an array literal',
+    builder.sourceLocation,
+  )
+  const slots = new Set<bigint>()
+  for (const item of builder.getItemBuilders()) {
+    if (item.resolvableToPType(numberPType)) {
+      slots.add(getLiteralNumber(item))
+    } else {
+      const from = getRangeProp(item, 'from')
+      const to = getRangeProp(item, 'to')
+      for (let i = from; i <= to; i++) {
+        slots.add(i)
+      }
+    }
+  }
+
+  return slots
+}
+
 function buildStateTotals(builder: NodeBuilder): ContractOptionsDecoratorData['stateTotals'] {
   function tryGetProp(name: string): bigint | undefined {
     if (builder.hasProperty(name)) {
-      return requireIntegerConstant(builder.memberAccess(name, builder.sourceLocation)).value
+      return getLiteralNumber(builder.memberAccess(name, builder.sourceLocation))
     }
     return undefined
   }

src/awst_build/eb/literal-expression-builder.ts

@@ -74,6 +74,6 @@ export abstract class LiteralExpressionBuilder extends InstanceBuilder {
   }
 
   hasProperty(_name: string): boolean {
-    this.throwInvalidExpression()
+    return false
   }
 }

src/awst_build/models/contract-class-model.ts

@@ -11,6 +11,7 @@ import { logger } from '../../logger'
 import type { Props } from '../../typescript-helpers'
 import { codeInvariant, invariant, isIn } from '../../util'
 import { CustomKeyMap } from '../../util/custom-key-map'
+import { intersectSets } from '../../util/intersect-sets'
 import type { ContractClassPType } from '../ptypes'
 import type { ContractOptionsDecoratorData } from './decorator-data'
 import { LogicSigClassModel } from './logic-sig-class-model'
@@ -58,14 +59,18 @@ export class ContractClassModel {
     let clearProgram: ContractMethod | null = this.clearProgram
     const methods: ContractMethod[] = [...this.methods, this.ctor]
     const methodResolutionOrder: ContractReference[] = []
-
     let firstBaseWithStateTotals: ContractClassModel | undefined = undefined
+    let reservedScratchSpace = new Set<bigint>()
+
     for (const baseType of this.type.allBases()) {
       const cref = ContractReference.fromPType(baseType)
       const baseClass = compilationSet.getContractClass(cref)
       if (baseClass.hasExplicitStateTotals() && firstBaseWithStateTotals === undefined) {
         firstBaseWithStateTotals = baseClass
       }
+      if (baseClass.options?.scratchSlots) {
+        reservedScratchSpace = intersectSets(reservedScratchSpace, baseClass.options.scratchSlots)
+      }
       methodResolutionOrder.push(cref)
       approvalProgram ??= baseClass.approvalProgram
       clearProgram ??= baseClass.clearProgram
@@ -114,18 +119,8 @@ export class ContractClassModel {
       localUints: this.options?.stateTotals?.localUints ?? null,
     })
 
-    // TODO: Tally from bases
-    const reservedScratchSpace = new Set<bigint>()
     if (this.options?.scratchSlots) {
-      for (const reservation of this.options.scratchSlots) {
-        if (typeof reservation === 'bigint') {
-          reservedScratchSpace.add(reservation)
-        } else {
-          for (let i = reservation.from; i <= reservation.to; i++) {
-            reservedScratchSpace.add(i)
-          }
-        }
-      }
+      reservedScratchSpace = intersectSets(reservedScratchSpace, this.options.scratchSlots)
     }
 
     return nodeFactory.contract({

src/awst_build/models/decorator-data.ts

@@ -35,13 +35,12 @@ export type LogicSigOptionsDecoratorData = {
   avmVersion?: SupportedAvmVersion
   name?: string
 }
-type NumberRange = { from: bigint; to: bigint }
 export type ContractOptionsDecoratorData = {
   type: typeof Constants.contractOptionsDecoratorName
   sourceLocation: SourceLocation
   avmVersion?: SupportedAvmVersion
   name?: string
-  scratchSlots?: Array<bigint | NumberRange>
+  scratchSlots?: Set<bigint>
   stateTotals?: {
     globalUints?: bigint
     globalBytes?: bigint

src/util/intersect-sets.ts

@@ -0,0 +1,3 @@
+export function intersectSets<T>(...sets: Set<T>[]): Set<T> {
+  return new Set<T>(sets.flatMap((s) => [...s]))
+}

tests/approvals/out/reserve-scratch/ReserveScratchAlgo.approval.teal

@@ -0,0 +1,70 @@
+#pragma version 10
+
+tests/approvals/reserve-scratch.algo.ts::ReserveScratchAlgo.approvalProgram:
+    intcblock 1 0 15 45
+    bytecblock "hello"
+    txn ApplicationID
+    bnz main_after_if_else@2
+    callsub constructor
+
+main_after_if_else@2:
+    // tests/approvals/reserve-scratch.algo.ts:13
+    // this.setThings()
+    callsub setThings
+    // tests/approvals/reserve-scratch.algo.ts:15
+    // assert(Scratch.loadUint64(0) === 1)
+    intc_1 // 0
+    loads
+    intc_0 // 1
+    ==
+    assert
+    // tests/approvals/reserve-scratch.algo.ts:16
+    // assert(Scratch.loadBytes(0) === Bytes('hello'))
+    intc_1 // 0
+    loads
+    bytec_0 // "hello"
+    ==
+    assert
+    // tests/approvals/reserve-scratch.algo.ts:17
+    // assert(Scratch.loadUint64(15) === 45)
+    intc_2 // 15
+    loads
+    intc_3 // 45
+    ==
+    assert
+    // tests/approvals/reserve-scratch.algo.ts:18
+    // return true
+    intc_0 // 1
+    return
+
+
+// tests/approvals/reserve-scratch.algo.ts::ReserveScratchAlgo.constructor() -> void:
+constructor:
+    // tests/approvals/reserve-scratch.algo.ts:4-5
+    // @contract({ scratchSlots: [0, 1, { from: 10, to: 20 }] })
+    // export class ReserveScratchAlgo extends BaseContract {
+    proto 0 0
+    retsub
+
+
+// tests/approvals/reserve-scratch.algo.ts::ReserveScratchAlgo.setThings() -> void:
+setThings:
+    // tests/approvals/reserve-scratch.algo.ts:6
+    // setThings() {
+    proto 0 0
+    // tests/approvals/reserve-scratch.algo.ts:7
+    // Scratch.store(0, 1)
+    intc_1 // 0
+    intc_0 // 1
+    stores
+    // tests/approvals/reserve-scratch.algo.ts:8
+    // Scratch.store(1, Bytes('hello'))
+    intc_0 // 1
+    bytec_0 // "hello"
+    stores
+    // tests/approvals/reserve-scratch.algo.ts:9
+    // Scratch.store(15, 45)
+    intc_2 // 15
+    intc_3 // 45
+    stores
+    retsub

tests/approvals/out/reserve-scratch/ReserveScratchAlgo.clear.teal

@@ -0,0 +1,5 @@
+#pragma version 10
+
+tests/approvals/reserve-scratch.algo.ts::ReserveScratchAlgo.clearStateProgram:
+    pushint 1 // 1
+    return

tests/approvals/out/reserve-scratch/ReserveScratchAlgo.ssa.ir

@@ -0,0 +1,37 @@
+contract tests/approvals/reserve-scratch.algo.ts::ReserveScratchAlgo:
+    program approval:
+        subroutine tests/approvals/reserve-scratch.algo.ts::ReserveScratchAlgo.approvalProgram() -> bool:
+            block@0: // L12
+                let reinterpret_bool%0#0: bool = (txn ApplicationID)
+                goto reinterpret_bool%0#0 ? block@2 : block@1
+            block@1: // if_body_L1
+                tests/approvals/reserve-scratch.algo.ts::ReserveScratchAlgo.constructor()
+                goto block@2
+            block@2: // after_if_else_L1
+                tests/approvals/reserve-scratch.algo.ts::ReserveScratchAlgo.setThings()
+                let tmp%0#0: uint64 = (loads 0u)
+                let tmp%1#0: bool = (== tmp%0#0 1u)
+                (assert tmp%1#0)
+                let tmp%2#0: bytes = (loads 0u)
+                let tmp%3#0: bool = (== tmp%2#0 "hello")
+                (assert tmp%3#0)
+                let tmp%4#0: uint64 = (loads 15u)
+                let tmp%5#0: bool = (== tmp%4#0 45u)
+                (assert tmp%5#0)
+                return 1u
+
+        subroutine tests/approvals/reserve-scratch.algo.ts::ReserveScratchAlgo.constructor() -> void:
+            block@0: // L4
+                return 
+
+        subroutine tests/approvals/reserve-scratch.algo.ts::ReserveScratchAlgo.setThings() -> void:
+            block@0: // L6
+                (stores 0u 1u)
+                (stores 1u "hello")
+                (stores 15u 45u)
+                return 
+
+    program clear-state:
+        subroutine @algorandfoundation/algorand-typescript/base-contract.d.ts::BaseContract.clearStateProgram() -> bool:
+            block@0: // L1
+                return 1u

tests/approvals/out/reserve-scratch/SubReserveScratchAlgo.approval.teal

@@ -0,0 +1,105 @@
+#pragma version 10
+
+tests/approvals/reserve-scratch.algo.ts::SubReserveScratchAlgo.approvalProgram:
+    intcblock 1 0 15 45
+    bytecblock "world" "hello"
+    txn ApplicationID
+    bnz main_after_if_else@2
+    callsub constructor
+
+main_after_if_else@2:
+    // tests/approvals/reserve-scratch.algo.ts:25
+    // super.approvalProgram()
+    callsub approvalProgram
+    pop
+    // tests/approvals/reserve-scratch.algo.ts:26
+    // Scratch.store(50, Bytes('world'))
+    pushint 50 // 50
+    bytec_0 // "world"
+    stores
+    // tests/approvals/reserve-scratch.algo.ts:27
+    // Scratch.store(16, Bytes('world'))
+    pushint 16 // 16
+    bytec_0 // "world"
+    stores
+    // tests/approvals/reserve-scratch.algo.ts:28
+    // return true
+    intc_0 // 1
+    return
+
+
+// tests/approvals/reserve-scratch.algo.ts::SubReserveScratchAlgo.constructor() -> void:
+constructor:
+    // tests/approvals/reserve-scratch.algo.ts:22-23
+    // @contract({ scratchSlots: [50] })
+    // export class SubReserveScratchAlgo extends ReserveScratchAlgo {
+    proto 0 0
+    callsub tests/approvals/reserve-scratch.algo.ts::ReserveScratchAlgo.constructor
+    retsub
+
+
+// tests/approvals/reserve-scratch.algo.ts::ReserveScratchAlgo.constructor() -> void:
+tests/approvals/reserve-scratch.algo.ts::ReserveScratchAlgo.constructor:
+    // tests/approvals/reserve-scratch.algo.ts:4-5
+    // @contract({ scratchSlots: [0, 1, { from: 10, to: 20 }] })
+    // export class ReserveScratchAlgo extends BaseContract {
+    proto 0 0
+    retsub
+
+
+// tests/approvals/reserve-scratch.algo.ts::ReserveScratchAlgo.approvalProgram() -> uint64:
+approvalProgram:
+    // tests/approvals/reserve-scratch.algo.ts:12
+    // approvalProgram(): boolean {
+    proto 0 1
+    // tests/approvals/reserve-scratch.algo.ts:13
+    // this.setThings()
+    callsub setThings
+    // tests/approvals/reserve-scratch.algo.ts:15
+    // assert(Scratch.loadUint64(0) === 1)
+    intc_1 // 0
+    loads
+    intc_0 // 1
+    ==
+    assert
+    // tests/approvals/reserve-scratch.algo.ts:16
+    // assert(Scratch.loadBytes(0) === Bytes('hello'))
+    intc_1 // 0
+    loads
+    bytec_1 // "hello"
+    ==
+    assert
+    // tests/approvals/reserve-scratch.algo.ts:17
+    // assert(Scratch.loadUint64(15) === 45)
+    intc_2 // 15
+    loads
+    intc_3 // 45
+    ==
+    assert
+    // tests/approvals/reserve-scratch.algo.ts:18
+    // return true
+    intc_0 // 1
+    retsub
+
+
+// tests/approvals/reserve-scratch.algo.ts::ReserveScratchAlgo.setThings() -> void:
+setThings:
+    // tests/approvals/reserve-scratch.algo.ts:6
+    // setThings() {
+    proto 0 0
+    // tests/approvals/reserve-scratch.algo.ts:7
+    // Scratch.store(0, 1)
+    intc_1 // 0
+    intc_0 // 1
+    stores
+    // tests/approvals/reserve-scratch.algo.ts:8
+    // Scratch.store(1, Bytes('hello'))
+    intc_0 // 1
+    bytec_1 // "hello"
+    stores
+    // tests/approvals/reserve-scratch.algo.ts:9
+    // Scratch.store(15, 45)
+    intc_2 // 15
+    intc_3 // 45
+    stores
+    retsub